def _fixmodel(arch, block, layers, pretrained, progress, **kwargs): model = ResNet(block, layers, **kwargs) pretrained_dict = load_state_dict_from_url(model_urls[arch], progress=progress, map_location='cpu')['model'] model_dict = model.state_dict() count = 0 count2 = 0 for k in model_dict.keys(): count = count + 1.0 if (('module.' + k) in pretrained_dict.keys()): count2 = count2 + 1.0 model_dict[k] = pretrained_dict.get(('module.' + k)) assert int(count2 * 100 / count) == 100, "model loading error" model.load_state_dict(model_dict) return model
def main(opt): start_epoch = 0 acc_best = 0. glob_step = 0 lr_now = opt.lr # save options log.save_options(opt, opt.ckpt) tb_logdir = f'./exp/{opt.name}' if os.path.exists(tb_logdir): shutil.rmtree(tb_logdir) writer = SummaryWriter(log_dir=f'./exp/{opt.name}') exp_dir_ = dirname(opt.load) # create model print(">>> creating model") # TODO: This is how to avoid weird data reshaping for non-3-channel inputs. # Have ResNet model take in grayscale rather than RGB # model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False) if opt.arch == 'cnn': model = ResNet(BasicBlock, [2, 2, 2, 2], num_classes=opt.num_classes) else: model = LinearModel() model = model.cuda() model.apply(weight_init) print(">>> total params: {:.2f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) criterion = nn.CrossEntropyLoss().cuda() optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr) # load ckpt if opt.load: print(">>> loading ckpt from '{}'".format(opt.load)) ckpt = torch.load(opt.load) start_epoch = ckpt['epoch'] acc_best = ckpt['acc'] glob_step = ckpt['step'] lr_now = ckpt['lr'] model.load_state_dict(ckpt['state_dict']) optimizer.load_state_dict(ckpt['optimizer']) print(">>> ckpt loaded (epoch: {} | acc: {})".format( start_epoch, acc_best)) if opt.resume: logger = log.Logger(os.path.join(opt.ckpt, 'log.txt'), resume=True) else: logger = log.Logger(os.path.join(opt.ckpt, 'log.txt')) logger.set_names([ 'epoch', 'lr', 'loss_train', 'err_train', 'acc_train', 'loss_test', 'err_test', 'acc_test' ]) transforms = [ ToTensor(), ] train_datasets = [] for dataset_name in opt.train_datasets: train_datasets.append( ClassificationDataset(name=dataset_name, num_kpts=opt.num_kpts, transforms=transforms, split='train', arch=opt.arch, gt=opt.gt)) train_dataset = ConcatDataset(train_datasets) train_loader = DataLoader(train_dataset, batch_size=opt.train_batch, shuffle=True, num_workers=opt.job) split = 'test' if opt.test else 'valid' test_dataset = ClassificationDataset(name=opt.test_dataset, num_kpts=opt.num_kpts, transforms=transforms, split=split, arch=opt.arch, gt=opt.gt) test_loader = DataLoader(test_dataset, batch_size=opt.test_batch, shuffle=False, num_workers=opt.job) subset_loaders = {} for subset in test_dataset.create_subsets(): subset_loaders[subset.split] = DataLoader(subset, batch_size=opt.test_batch, shuffle=False, num_workers=opt.job) cudnn.benchmark = True for epoch in range(start_epoch, opt.epochs): torch.cuda.empty_cache() print('==========================') print('>>> epoch: {} | lr: {:.5f}'.format(epoch + 1, lr_now)) if not opt.test: glob_step, lr_now, loss_train, err_train, acc_train = \ train(train_loader, model, criterion, optimizer, num_kpts=opt.num_kpts, num_classes=opt.num_classes, lr_init=opt.lr, lr_now=lr_now, glob_step=glob_step, lr_decay=opt.lr_decay, gamma=opt.lr_gamma, max_norm=opt.max_norm) loss_test, err_test, acc_test, auc_test, prec_test = \ test(test_loader, model, criterion, num_kpts=opt.num_kpts, num_classes=opt.num_classes, batch_size=opt.test_batch) ## Test subsets ## subset_losses = {} subset_errs = {} subset_accs = {} subset_aucs = {} subset_precs = {} subset_openpose = {} subset_missing = {} subset_grids = {} if len(subset_loaders) > 0: bar = Bar('>>>', fill='>', max=len(subset_loaders)) for key_idx, key in enumerate(subset_loaders): loss_sub, err_sub, acc_sub, auc_sub, prec_sub = test( subset_loaders[key], model, criterion, num_kpts=opt.num_kpts, num_classes=opt.num_classes, batch_size=4, log=False) subset_losses[key] = loss_sub subset_errs[key] = err_sub subset_accs[key] = acc_sub subset_aucs[key] = auc_sub subset_precs[key] = prec_sub sub_dataset = subset_loaders[key].dataset if sub_dataset.gt_paths is not None: gt_X = load_gt(sub_dataset.gt_paths) subset_openpose[key] = mpjpe_2d_openpose(sub_dataset.X, gt_X) subset_missing[key] = mean_missing_parts(sub_dataset.X) else: subset_openpose[key] = 0. subset_missing[key] = 0. sample_idxs = extract_tb_sample(subset_loaders[key], model, batch_size=opt.test_batch) sample_X = sub_dataset.X[sample_idxs] sample_img_paths = [sub_dataset.img_paths[x] for x in sample_idxs] if opt.arch == 'cnn': subset_grids[key] = create_grid(sample_X, sample_img_paths) bar.suffix = f'({key_idx+1}/{len(subset_loaders)}) | {key}' bar.next() if len(subset_loaders) > 0: bar.finish() ################### if opt.test: subset_accs['all'] = acc_test subset_aucs['all'] = auc_test subset_precs['all'] = prec_test report_dict = { 'acc': subset_accs, 'auc': subset_aucs, 'prec': subset_precs } report_idx = 0 report_path = f'report/{opt.name}-{report_idx}.json' while os.path.exists(f'report/{opt.name}-{report_idx}.json'): report_idx += 1 report_path = f'report/{opt.name}-{report_idx}.json' print(f'>>> Saving report to {report_path}...') with open(report_path, 'w') as acc_f: json.dump(report_dict, acc_f, indent=4) print('>>> Exiting (test mode)...') break # update log file logger.append([ epoch + 1, lr_now, loss_train, err_train, acc_train, loss_test, err_test, acc_test ], [ 'int', 'float', 'float', 'float', 'float', 'float', 'float', 'float' ]) # save ckpt is_best = acc_test > acc_best acc_best = max(acc_test, acc_best) if is_best: log.save_ckpt( { 'epoch': epoch + 1, 'lr': lr_now, 'step': glob_step, 'acc': acc_best, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict() }, ckpt_path=opt.ckpt, is_best=True) else: log.save_ckpt( { 'epoch': epoch + 1, 'lr': lr_now, 'step': glob_step, 'acc': acc_best, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict() }, ckpt_path=opt.ckpt, is_best=False) writer.add_scalar('Loss/train', loss_train, epoch) writer.add_scalar('Loss/test', loss_test, epoch) writer.add_scalar('Error/train', err_train, epoch) writer.add_scalar('Error/test', err_test, epoch) writer.add_scalar('Accuracy/train', acc_train, epoch) writer.add_scalar('Accuracy/test', acc_test, epoch) for key in subset_losses: writer.add_scalar(f'Loss/Subsets/{key}', subset_losses[key], epoch) writer.add_scalar(f'Error/Subsets/{key}', subset_errs[key], epoch) writer.add_scalar(f'Accuracy/Subsets/{key}', subset_accs[key], epoch) writer.add_scalar(f'OpenPose/Subsets/{key}', subset_openpose[key], epoch) writer.add_scalar(f'Missing/Subsets/{key}', subset_missing[key], epoch) if opt.arch == 'cnn': writer.add_images(f'Subsets/{key}', subset_grids[key], epoch, dataformats='NHWC') logger.close() writer.close()