| import random | |
| import argparse | |
| import yaml | |
| import torch | |
| import os | |
| import util | |
| from util import build_model, train_one_epoch | |
| from dataloader2 import generate_dataset_loader | |
| import torch.nn as nn | |
| from torch.optim.lr_scheduler import StepLR, LambdaLR, MultiStepLR | |
| def get_arguments(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--config', dest='config', help='settings of detector in yaml format') | |
| args = parser.parse_args() | |
| return args | |
| if __name__ == '__main__': | |
| args = get_arguments() | |
| assert (os.path.exists(args.config)) | |
| cfg = yaml.load(open(args.config, 'r'), Loader=yaml.Loader) | |
| print("******* Building models. *******") | |
| print(cfg) | |
| model = util.build_model(cfg['model']) | |
| model = model.cuda() | |
| if cfg['tuning_mode'] == 'lp': | |
| for param in model.encoder.parameters(): | |
| param.requires_grad = False | |
| model = torch.nn.DataParallel(model) | |
| optimizer = torch.optim.AdamW(model.parameters(), lr=cfg['lr'], weight_decay=1e-8) | |
| scheduler = MultiStepLR(optimizer, milestones=[20, 25], gamma=0.1) | |
| loss = nn.BCEWithLogitsLoss() | |
| trMaxEpoch = cfg['max_epoch'] | |
| snapshot_path = cfg['save_dir'] | |
| if not os.path.exists(snapshot_path): | |
| os.makedirs(snapshot_path) | |
| max_epoch, max_acc = 0, 0 | |
| for epochID in range(0, trMaxEpoch): | |
| print("******* Training epoch", str(epochID)," *******") | |
| print("******* Building datasets. *******") | |
| train_loader, val_loader = generate_dataset_loader(cfg) | |
| max_epoch, max_acc, epoch_time = train_one_epoch(cfg, model, loss, scheduler, optimizer, epochID, max_epoch, max_acc, train_loader, val_loader, snapshot_path) | |
| print("******* Ending epoch", str(epochID)," Time ", str(epoch_time), "*******") | |