295 lines
10 KiB
Python
Executable File
295 lines
10 KiB
Python
Executable File
import argparse
|
|
import os
|
|
import os.path as osp
|
|
|
|
import torch
|
|
import wandb
|
|
import yaml
|
|
from criterions import build_criterion
|
|
from model import build_model
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
from trainer import build_eval, build_trainer
|
|
from utils import *
|
|
from utils import get_logger
|
|
|
|
from datasets import build_data_loader
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", type=str)
|
|
parser.add_argument("--eval", action="store_true")
|
|
parser.add_argument("--ckpt", type=str)
|
|
parser.add_argument("--amp", action="store_true")
|
|
parser.add_argument("--tensorboard", action="store_true")
|
|
parser.add_argument("--wandb", action="store_true")
|
|
parser.add_argument("--lr_scheduler", action="store_true")
|
|
parser.add_argument("--no_rgb", action="store_true")
|
|
parser.add_argument("--no_flow", action="store_true")
|
|
parser.add_argument(
|
|
"--save-as-cfg",
|
|
action="store_true",
|
|
help="Save outputs with config name in the output directory",
|
|
)
|
|
parser.add_argument("--debug", action="store_true")
|
|
parser.add_argument(
|
|
"--save-preds",
|
|
action="store_true",
|
|
help="Save predictions in the output directory",
|
|
)
|
|
# parser.add_argument('--model', type=str, default=None)
|
|
args = parser.parse_args()
|
|
|
|
# combine argparse and yaml
|
|
opt = yaml.load(open(args.config), Loader=yaml.FullLoader)
|
|
opt.update(vars(args))
|
|
cfg = opt
|
|
|
|
if "checkpoint" in cfg and cfg["checkpoint"] is not None and args.ckpt is None:
|
|
args.ckpt = cfg["checkpoint"]
|
|
print(f"checkpoint is given in the config, set args.ckpt to {args.ckpt}")
|
|
|
|
seed = 20 if "seed" not in cfg else cfg["seed"]
|
|
set_seed(seed)
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
annotation_type = os.path.basename(cfg["annotation_type"])
|
|
|
|
identifier = os.path.basename(cfg["config"])[:-5]
|
|
identifier += "_"
|
|
if cfg["no_flow"] == False:
|
|
identifier += "A"
|
|
if cfg["no_rgb"] == False:
|
|
identifier += "V"
|
|
|
|
if cfg["save_as_cfg"] == True:
|
|
output_dir = osp.join("output_cfg", identifier)
|
|
result_path = create_outdir(output_dir, overwrite=True)
|
|
else:
|
|
identifier = get_identifier(cfg)
|
|
annotation_type = os.path.basename(cfg["annotation_type"])
|
|
identifier += f"_{annotation_type}"
|
|
result_path = create_outdir(osp.join(cfg["output_path"], identifier))
|
|
|
|
logger = get_logger(result_path)
|
|
logger.info(f"Identifier: {identifier}")
|
|
logger.info(f"result_path: {result_path}")
|
|
logger.info(cfg)
|
|
|
|
testloader = build_data_loader(cfg, mode="test")
|
|
print("testloader loaded")
|
|
model = build_model(cfg, device)
|
|
evaluate = build_eval(cfg)
|
|
|
|
trainloader = build_data_loader(cfg, mode="train")
|
|
print("trainloader loaded")
|
|
|
|
# evaluation
|
|
if args.eval == True and cfg["ckpt"] == None:
|
|
print("Automatically find the best checkpoint")
|
|
ckpt_root = os.path.join("output_cfg", identifier, "ckpts")
|
|
ckpt_dirs = os.listdir(ckpt_root)
|
|
|
|
ckpt_path = os.path.join(ckpt_root, "best.pth")
|
|
if os.path.exists(ckpt_path):
|
|
args.ckpt = ckpt_path
|
|
print(f"args.ckpt set to {args.ckpt}")
|
|
|
|
else:
|
|
ckpt_path = os.path.join(ckpt_root, "pred", "best.pth")
|
|
if os.path.exists(ckpt_path):
|
|
args.ckpt = ckpt_path
|
|
print(f"args.ckpt set to {args.ckpt}")
|
|
|
|
print("Warninig) no ckpt_path found, there can be multiple ckpts")
|
|
print(f"ckpt_dirs: {ckpt_dirs}")
|
|
pred_dir_number = 0
|
|
ckpt_dirs.sort(reverse=True)
|
|
|
|
if args.ckpt == None:
|
|
for ckpt_dir in ckpt_dirs:
|
|
ckpt_path = os.path.join(ckpt_root, ckpt_dir, "best.pth")
|
|
|
|
if os.path.exists(ckpt_path):
|
|
args.ckpt = ckpt_path
|
|
print(f"Found ckpt, args.ckpt was None, set to {args.ckpt}")
|
|
break
|
|
|
|
if args.ckpt == None:
|
|
raise ValueError("No checkpoint found, args.ckpt is None")
|
|
|
|
# load checkpoint if args.ckpt is given
|
|
if args.eval == True:
|
|
if args.ckpt == None:
|
|
raise ValueError("Please specify the checkpoint path")
|
|
checkpoint = torch.load(args.ckpt)
|
|
model.load_state_dict(checkpoint)
|
|
|
|
elif args.eval == False and args.ckpt != None:
|
|
checkpoint = torch.load(args.ckpt)
|
|
print(f"Load checkpoint from {args.ckpt}")
|
|
model_dict = model.state_dict()
|
|
|
|
unnecessary_keys = []
|
|
matched_ckpt = {}
|
|
# Filter out unnecessary keys
|
|
for k, v in checkpoint.items():
|
|
if not (k in model_dict and model_dict[k].shape == checkpoint[k].shape):
|
|
unnecessary_keys.append(k)
|
|
elif (
|
|
"except_classifier" in cfg
|
|
and cfg["except_classifier"] == True
|
|
and "classification" in k
|
|
):
|
|
print(f"Unnecessary key: {k} is classifier")
|
|
unnecessary_keys.append(k)
|
|
else:
|
|
matched_ckpt[k] = v
|
|
|
|
# Overwrite entries in the existing state dict
|
|
model_dict.update(matched_ckpt)
|
|
|
|
# Load the new state dict
|
|
model.load_state_dict(model_dict)
|
|
|
|
print(f"Unnecessary keys: {unnecessary_keys}")
|
|
|
|
if cfg["wandb"] == True:
|
|
# args.tensorboard = True
|
|
print("wandb init")
|
|
os.environ["WANDB__SERVICE_WAIT"] = "3000"
|
|
wandb.init(
|
|
project="egospeak",
|
|
tags=[cfg["model"], cfg["data_name"], annotation_type],
|
|
config=cfg,
|
|
name=os.path.basename(result_path),
|
|
)
|
|
|
|
if args.eval == True or args.debug == True:
|
|
if args.eval == True and args.ckpt == None:
|
|
raise ValueError("Please specify the checkpoint path")
|
|
|
|
logger.info(f'Dataset: {cfg["data_name"]}, Model: {cfg["model"]}')
|
|
logger.info(f"Result path:{result_path}")
|
|
|
|
for phase in ["test"]:
|
|
if phase == "train":
|
|
raise NotImplementedError("evaluation on train split is not implemented")
|
|
dataloader = trainloader
|
|
|
|
elif phase == "test":
|
|
dataloader = testloader
|
|
|
|
mAP, mcAP, preds = evaluate(
|
|
model, dataloader, logger, result_path=result_path
|
|
)
|
|
|
|
pred_path = osp.join(result_path, "eval")
|
|
os.makedirs(pred_path, exist_ok=True)
|
|
|
|
if args.save_preds == True:
|
|
print(f"pred_path: {pred_path}")
|
|
for k, v in preds.items():
|
|
np.save(osp.join(pred_path, f"{k}.npy"), v)
|
|
|
|
logger.info(f'{phase}) {cfg["task"]} result: {mAP*100:.2f} mAP')
|
|
logger.info(f'{phase}) {cfg["task"]} result: {mcAP*100:.2f} mcAP')
|
|
exit()
|
|
|
|
criterion = build_criterion(cfg, device)
|
|
train_one_epoch = build_trainer(cfg)
|
|
optim = torch.optim.AdamW if cfg["optimizer"] == "AdamW" else torch.optim.Adam
|
|
optimizer = optim(
|
|
[{"params": model.parameters(), "initial_lr": cfg["lr"]}],
|
|
lr=cfg["lr"],
|
|
weight_decay=cfg["weight_decay"],
|
|
)
|
|
|
|
scheduler = (
|
|
build_lr_scheduler(cfg, optimizer, len(trainloader))
|
|
if args.lr_scheduler
|
|
else None
|
|
)
|
|
writer = SummaryWriter(osp.join(result_path, "runs")) if args.tensorboard else None
|
|
scaler = torch.cuda.amp.GradScaler() if args.amp else None
|
|
total_params = sum(p.numel() for p in model.parameters())
|
|
|
|
logger.info(f'Dataset: {cfg["data_name"]}, Model: {cfg["model"]}')
|
|
logger.info(
|
|
f'lr:{cfg["lr"]} | Weight Decay:{cfg["weight_decay"]} | Window Size:{cfg["window_size"]} | Batch Size:{cfg["batch_size"]}'
|
|
)
|
|
logger.info(
|
|
f'Total epoch:{cfg["num_epoch"]} | Total Params:{total_params/1e6:.1f} M | Optimizer: {cfg["optimizer"]}'
|
|
)
|
|
logger.info(f"Output Path:{result_path}")
|
|
|
|
best_mAP, best_epoch = 0, 0
|
|
best_cAP, best_cAP_epoch = 0, 0
|
|
|
|
pred_save_path = osp.join(result_path, "ckpts", "pred")
|
|
print(f"pred_save_path: {pred_save_path}")
|
|
|
|
# if you want to overwrite, change the code
|
|
os.makedirs(pred_save_path, exist_ok=True)
|
|
|
|
for epoch in range(1, cfg["num_epoch"] + 1):
|
|
epoch_loss = train_one_epoch(
|
|
trainloader,
|
|
model,
|
|
criterion,
|
|
optimizer,
|
|
scaler,
|
|
epoch,
|
|
cfg["wandb"],
|
|
writer,
|
|
scheduler=scheduler,
|
|
)
|
|
trainloader.dataset._init_features()
|
|
|
|
mAP, cAP, preds = evaluate(
|
|
model, testloader, logger, criterion, cfg["wandb"], epoch, writer
|
|
)
|
|
|
|
if mAP > best_mAP:
|
|
best_mAP = mAP
|
|
best_epoch = epoch
|
|
torch.save(model.state_dict(), osp.join(pred_save_path, "best.pth"))
|
|
|
|
if args.save_preds == True:
|
|
for k, v in preds.items():
|
|
np.save(osp.join(pred_save_path, f"{k}.npy"), v)
|
|
|
|
if cAP > best_cAP:
|
|
best_cAP = cAP
|
|
best_cAP_epoch = epoch
|
|
torch.save(model.state_dict(), osp.join(pred_save_path, "best_cAP.pth"))
|
|
|
|
logger.info(
|
|
f'Epoch {epoch} mAP: {mAP*100:.2f} | Best mAP: {best_mAP*100:.2f} at epoch {best_epoch}, iter {epoch*cfg["batch_size"]*len(trainloader)} | train_loss: {epoch_loss/len(trainloader):.4f}, lr: {optimizer.param_groups[0]["lr"]:.7f}'
|
|
)
|
|
logger.info(
|
|
f'Epoch {epoch} cAP: {cAP*100:.2f} | Best cAP: {best_cAP*100:.2f} at epoch {best_epoch}, iter {epoch*cfg["batch_size"]*len(trainloader)} | train_loss: {epoch_loss/len(trainloader):.4f}, lr: {optimizer.param_groups[0]["lr"]:.7f}'
|
|
)
|
|
|
|
|
|
if cfg["wandb"] == True:
|
|
logger.info("Wandb logging")
|
|
wandb.log(
|
|
{
|
|
"Epoch": epoch,
|
|
"Train Loss (Epoch)": epoch_loss / len(trainloader),
|
|
"Learning Rate": optimizer.param_groups[0]["lr"],
|
|
"anticipation average mAP": mAP,
|
|
"anticipation average cAP": cAP,
|
|
}
|
|
)
|
|
|
|
# os.rename(osp.join(result_path, 'ckpts', 'best.pth'), osp.join(result_path, 'ckpts', f'best_{best_mAP*100:.2f}.pth'))
|
|
# save the training result and config as txt file
|
|
with open(osp.join(result_path, "result.txt"), "w") as f:
|
|
f.write(f"Best mAP: {best_mAP*100:.2f} at epoch {best_epoch}\n")
|
|
f.write(f"Best cAP: {best_cAP*100:.2f} at epoch {best_cAP_epoch}\n")
|
|
f.write(f"Config: {cfg}")
|
|
|
|
wandb.finish() if args.wandb else None
|