def __init__(self, gallery_path=args.gallery_path, save_dir=args.save_dir): print("==========\nArgs:{}\n==========".format(args)) self.gallery_path = gallery_path self.save_dir=save_dir print("gallery_path".format(gallery_path)) print("Initializing image data manager") # ------------ self.galleryloader, self.gallery_datasets = self.create_galleryloader() print("Initializing model: {}".format(args.arch)) self.model = models.init_model(name=args.arch, num_classes=args.num_train_pids, loss={'xent', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(self.model))) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size map_location = 'cpu' if not use_gpu else None checkpoint = torch.load(args.load_weights, map_location=map_location) pretrain_dict = checkpoint['state_dict'] model_dict = self.model.state_dict() pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()} model_dict.update(pretrain_dict) self.model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if use_gpu: self.model = nn.DataParallel(self.model).cuda() self.gf = self.init_gallary_info()
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False logger_info = LoggerInfo() sys.stdout = Logger(logger_info) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("\nInitializing dataset {}".format(args.dataset_plt)) dataset_plt = data_manager.init_imgreid_dataset(root=args.root, name=args.dataset_plt) print("\nInitializing dataset {}".format(args.dataset_vecl)) dataset_vecl = data_manager.init_imgreid_dataset(root=args.root, name=args.dataset_vecl) transform_test_plt = T.Compose([ T.Resize((args.height_plt, args.width_plt)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # transform_flip_test_plt = T.Compose([ # T.Resize((args.height_plt, args.width_plt)), # functional.hflip, # T.ToTensor(), # T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # ]) transform_test_vecl = T.Compose([ T.Resize((args.height_vecl, args.width_vecl)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # transform_flip_test_vecl = T.Compose([ # T.Resize((args.height_vecl, args.width_vecl)), # functional.hflip, # T.ToTensor(), # T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # ]) pin_memory = True if use_gpu else False queryloader_plt = DataLoader( ImageDatasetV2(dataset_plt.query, transform=transform_test_plt), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) # queryloader_flip_plt = DataLoader( # ImageDatasetV2(dataset_plt.query, transform=transform_flip_test_plt), # batch_size=args.test_batch, shuffle=False, num_workers=args.workers, # pin_memory=pin_memory, drop_last=False, # ) # queryloader_plt = [queryloader_plt, queryloader_flip_plt] queryloader_plt = [queryloader_plt] galleryloader_plt = DataLoader( ImageDatasetV2(dataset_plt.gallery, transform=transform_test_plt), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) # galleryloader_flip_plt = DataLoader( # ImageDatasetV2(dataset_plt.gallery, transform=transform_flip_test_plt), # batch_size=args.test_batch, shuffle=False, num_workers=args.workers, # pin_memory=pin_memory, drop_last=False, # ) # galleryloader_plt = [galleryloader_plt, galleryloader_flip_plt] galleryloader_plt = [galleryloader_plt] queryloader_vecl = DataLoader( ImageDatasetWGL(dataset_vecl.query, transform=transform_test_vecl, with_image_name=True), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) # queryloader_flip_vecl = DataLoader( # ImageDatasetV2(dataset_vecl.query, transform=transform_flip_test_vecl), # batch_size=args.test_batch, shuffle=False, num_workers=args.workers, # pin_memory=pin_memory, drop_last=False, # ) # queryloader_vecl = [queryloader_vecl, queryloader_flip_vecl] queryloader_vecl = [queryloader_vecl] galleryloader_vecl = DataLoader( ImageDatasetWGL(dataset_vecl.gallery, transform=transform_test_vecl, with_image_name=True), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) # galleryloader_flip_vecl = DataLoader( # ImageDatasetV2(dataset_vecl.gallery, transform=transform_flip_test_vecl), # batch_size=args.test_batch, shuffle=False, num_workers=args.workers, # pin_memory=pin_memory, drop_last=False, # ) # galleryloader_vecl = [galleryloader_vecl, galleryloader_flip_vecl] galleryloader_vecl = [galleryloader_vecl] print("\nInitializing model: {}".format(args.arch)) model_plt = models.init_model(name=args.arch_plt, num_classes=dataset_plt.num_train_pids, loss_type=args.loss_type) model_vecl = models.init_model(name=args.arch_vecl, num_classes=dataset_vecl.num_train_pids, loss_type=args.loss_type) print("Plate model size: {:.3f} M".format(count_num_param(model_plt))) print("Vehicle model size: {:.3f} M".format(count_num_param(model_vecl))) if args.loss_type == 'xent': criterion = nn.CrossEntropyLoss() else: raise KeyError("Unsupported loss: {}".format(args.loss_type)) if args.resm_plt and args.resm_vecl: if check_isfile(args.resm_plt) and check_isfile(args.resm_vecl): ckpt_plt = torch.load(args.resm_plt) pre_dic_plt = ckpt_plt['state_dict'] model_dic_plt = model_plt.state_dict() pre_dic_plt = { k: v for k, v in pre_dic_plt.items() if k in model_dic_plt and model_dic_plt[k].size() == v.size() } model_dic_plt.update(pre_dic_plt) model_plt.load_state_dict(model_dic_plt) args.start_epoch_plt = ckpt_plt['epoch'] rank1_plt = ckpt_plt['rank1'] ckpt_vecl = torch.load(args.resm_vecl) pre_dic_vecl = ckpt_vecl['state_dict'] model_dic_vecl = model_vecl.state_dict() pre_dic_vecl = { k: v for k, v in pre_dic_vecl.items() if k in model_dic_vecl and model_dic_vecl[k].size() == v.size() } model_dic_vecl.update(pre_dic_vecl) model_vecl.load_state_dict(model_dic_vecl) args.start_epoch_vecl = ckpt_vecl['epoch'] rank1_vecl = ckpt_vecl['rank1'] print("\nLoaded checkpoint from '{}' \nand '{}".format( args.resm_plt, args.resm_vecl)) print("Plate model: start_epoch: {}, rank1: {}".format( args.start_epoch_plt, rank1_plt)) print("Vehicle model: start_epoch: {}, rank1: {}".format( args.start_epoch_vecl, rank1_vecl)) if use_gpu: model_plt = nn.DataParallel(model_plt).cuda() model_vecl = nn.DataParallel(model_vecl).cuda() if args.evaluate: print("\nEvaluate only") test(model_plt, model_vecl, queryloader_plt, queryloader_vecl, galleryloader_plt, galleryloader_vecl, use_gpu) return
def main(): global args torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'log_train_{}.txt'.format(time.strftime("%Y-%m-%d-%H-%M-%S")) if args.evaluate: log_name.replace('train', 'test') sys.stdout = Logger(osp.join(args.save_dir, log_name)) print(' '.join(sys.argv)) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU, however, GPU is highly recommended") print("Initializing image data manager") dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args)) if hasattr(dm, 'lfw_dataset'): lfw = dm.lfw_dataset print('LFW dataset is used!') else: lfw = None trainloader, trainloader_dict, testloader_dict = dm.return_dataloaders() num_train_pids = dm.num_train_pids print("Initializing model: {}".format(args.arch)) model = models.init_model( name=args.arch, num_classes=num_train_pids, loss={'xent', 'htri'}, pretrained=False if args.load_weights else 'imagenet', grayscale=args.grayscale, normalize_embeddings=args.normalize_embeddings, normalize_fc=args.normalize_fc, convbn=args.convbn) print("Model size: {:.3f} M".format(count_num_param(model))) count_flops(model, args.height, args.width, args.grayscale) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size load_weights(model, args.load_weights) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, checkpoint['rank1'])) if use_gpu: model = nn.DataParallel(model).cuda() model = model.cuda() if args.evaluate: print("Evaluate only") for name in args.target_names: if not 'lfw' in name.lower(): print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] distmat = test(args, model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results( distmat, dm.return_testdataset_by_name(name), save_dir=osp.join(args.save_dir, 'ranked_results', name), topk=20) else: model.eval() same_acc, diff_acc, all_acc, auc, thresh = evaluate( args, dm.lfw_dataset, model, compute_embeddings_lfw, args.test_batch_size, verbose=False, show_failed=args.show_failed) log.info('Validation accuracy: {0:.4f}, {1:.4f}'.format( same_acc, diff_acc)) log.info('Validation accuracy mean: {0:.4f}'.format(all_acc)) log.info('Validation AUC: {0:.4f}'.format(auc)) log.info('Estimated threshold: {0:.4f}'.format(thresh)) return criterions = choose_losses(args, dm, model, use_gpu) if not args.evaluate and len(criterions) == 0: raise AssertionError('No loss functions were chosen!') optimizer = init_optimizer(model.parameters(), **optimizer_kwargs(args)) if args.load_optim: checkpoint = torch.load(args.load_weights) optimizer.load_state_dict(checkpoint['optim']) print("Loaded optimizer from '{}'".format(args.load_weights)) for param_group in optimizer.param_groups: param_group['lr'] = args.lr param_group['weight_decay'] = args.weight_decay scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) start_time = time.time() ranklogger = RankLogger(args.source_names, args.target_names) train_time = 0 train_writer = SummaryWriter(osp.join(args.save_dir, 'train_log')) test_writer = SummaryWriter(osp.join(args.save_dir, 'test_log')) print("=> Start training") if args.fixbase_epoch > 0: print( "Train {} for {} epochs while keeping other layers frozen".format( args.open_layers, args.fixbase_epoch)) initial_optim_state = optimizer.state_dict() for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterions, optimizer, trainloader, use_gpu, train_writer, fixbase=True, lfw=lfw) train_time += round(time.time() - start_train_time) for name in args.target_names: if not 'lfw' in name.lower(): print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] testloader = testloader_dict[name]['test'] criteria = None rank1 = test(args, model, queryloader, galleryloader, use_gpu, testloader=testloader, criterions=criteria) else: model.eval() same_acc, diff_acc, all_acc, auc, thresh = evaluate( args, dm.lfw_dataset, model, compute_embeddings_lfw, args.test_batch_size, verbose=False, show_failed=args.show_failed) print('Validation accuracy: {0:.4f}, {1:.4f}'.format( same_acc, diff_acc)) print('Validation accuracy mean: {0:.4f}'.format(all_acc)) print('Validation AUC: {0:.4f}'.format(auc)) print('Estimated threshold: {0:.4f}'.format(thresh)) rank1 = all_acc print("Done. All layers are open to train for {} epochs".format( args.max_epoch)) optimizer.load_state_dict(initial_optim_state) for epoch in range(args.start_epoch, args.max_epoch): for criterion in criterions: criterion.train_stats.reset() start_train_time = time.time() train(epoch, model, criterions, optimizer, trainloader, use_gpu, train_writer, lfw=lfw) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_freq > 0 and ( epoch + 1) % args.eval_freq == 0 or (epoch + 1) == args.max_epoch: num_iter = (epoch + 1) * len(trainloader) if not args.no_train_quality: for name in args.source_names: print( "Measure quality on the {} train set...".format(name)) queryloader = trainloader_dict[name]['query'] galleryloader = trainloader_dict[name]['gallery'] rank1 = test(args, model, queryloader, galleryloader, use_gpu) train_writer.add_scalar('rank1/{}'.format(name), rank1, num_iter) print("=> Test") for name in args.target_names: if not 'lfw' in name.lower(): print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] testloader = testloader_dict[name]['test'] criteria = criterions if args.no_loss_on_val: criteria = None rank1 = test(args, model, queryloader, galleryloader, use_gpu, testloader=testloader, criterions=criteria) test_writer.add_scalar('rank1/{}'.format(name), rank1, num_iter) if not args.no_loss_on_val: for criterion in criterions: test_writer.add_scalar( 'loss/{}'.format(criterion.name), criterion.test_stats.avg, num_iter) criterion.test_stats.reset() ranklogger.write(name, epoch + 1, rank1) else: model.eval() same_acc, diff_acc, all_acc, auc, thresh = evaluate( args, dm.lfw_dataset, model, compute_embeddings_lfw, args.test_batch_size, verbose=False, show_failed=args.show_failed) print('Validation accuracy: {0:.4f}, {1:.4f}'.format( same_acc, diff_acc)) print('Validation accuracy mean: {0:.4f}'.format(all_acc)) print('Validation AUC: {0:.4f}'.format(auc)) print('Estimated threshold: {0:.4f}'.format(thresh)) test_writer.add_scalar('Accuracy/Val_same_accuracy', same_acc, num_iter) test_writer.add_scalar('Accuracy/Val_diff_accuracy', diff_acc, num_iter) test_writer.add_scalar('Accuracy/Val_accuracy', all_acc, num_iter) test_writer.add_scalar('Accuracy/AUC', auc, num_iter) rank1 = all_acc if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_dict = { 'state_dict': state_dict, 'epoch': epoch, 'optim': optimizer.state_dict() } if len(args.target_names): save_dict['rank1'] = rank1 save_checkpoint( save_dict, False, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time)) ranklogger.show_summary()
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset( root=args.root, name=args.dataset, split_id=args.split_id, cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split, ) #cuhk03_labeled: detected,labeled transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False #pdb.set_trace() trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), sampler=RandomIdentitySampler(dataset.train, args.train_batch, args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( ImageDataset(dataset.query, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( ImageDataset(dataset.gallery, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, batchsize=args.test_batch, loss={'xent', 'wcont', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(model))) criterion_xent = nn.CrossEntropyLoss() criterion_htri = TripletLoss(margin=args.margin) criterion_KA = KALoss(margin=args.margin, same_margin=args.same_margin, use_auto_samemargin=args.use_auto_samemargin) cirterion_lifted = LiftedLoss(margin=args.margin) cirterion_batri = BA_TripletLoss(margin=args.margin) if args.use_auto_samemargin == True: G_params = [{ 'params': model.parameters(), 'lr': args.lr }, { 'params': criterion_KA.auto_samemargin, 'lr': args.lr }] else: G_params = [para for _, para in model.named_parameters()] optimizer = init_optim(args.optim, G_params, args.lr, args.weight_decay) if args.load_weights: # load pretrained weights but ignore layers that don't match in size if check_isfile(args.load_weights): checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format( args.load_weights)) if args.resume: if check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format( args.start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.vis_ranked_res: visualize_ranked_results( distmat, dataset, save_dir=osp.join(args.save_dir, 'ranked_results'), topk=20, ) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() adjust_learning_rate(optimizer, epoch) train(epoch, model, cirterion_batri, cirterion_lifted, criterion_xent, criterion_htri, criterion_KA, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: rank1 = 0 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() print("==> Test") sys.stdout.flush() rank1 = test(model, queryloader, galleryloader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("model saved") print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) sys.stdout.flush() elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time)) sys.stdout.flush()
def main(): parser = argparse.ArgumentParser() parser.add_argument( '--snap_shot', type=str, default='saved-models/densenet121_xent_market1501.pth.tar') parser.add_argument('--arch', type=str, default='densenet121') parser.add_argument('--dataset-path', type=str, default='data/valset/valSet') parser.add_argument('--height', type=int, default=256, help="height of an image (default: 256)") parser.add_argument('--width', type=int, default=128, help="width of an image (default: 128)") parser.add_argument('--test-batch', default=100, type=int, help="test batch size") parser.add_argument('-j', '--workers', default=4, type=int, help="number of data loading workers (default: 4)") parser.add_argument('--log-dir', type=str, default='log/eval_625') parser.add_argument('--gpu', type=int, default=1) args = parser.parse_args() pin_memory = True if args.gpu else False print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=751, loss={'xent'}, use_gpu=args.gpu).cuda() print("Model size: {:.3f} M".format(count_num_param(model))) checkpoint = torch.load(args.snap_shot) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.snap_shot)) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) queryloader = DataLoader( evalDataset(os.path.join(args.dataset_path, 'query'), transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( evalDataset(os.path.join(args.dataset_path, 'gallery'), transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) dataloaders = {'query': queryloader, 'gallery': galleryloader} for dataset in ['val']: for subset in ['query', 'gallery']: test_names, test_features = extractor(model, dataloaders[subset]) results = {'names': test_names, 'features': test_features.numpy()} scipy.io.savemat( os.path.join(args.log_dir, 'feature_%s_%s.mat' % (dataset, subset)), results)
def main(): global args set_random_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'log_test.txt' if args.evaluate else 'log_train.txt' sys.stdout = Logger(osp.join(args.save_dir, log_name)) print('==========\nArgs:{}\n=========='.format(args)) if use_gpu: print('Currently using GPU {}'.format(args.gpu_devices)) cudnn.benchmark = True else: print('Currently using CPU, however, GPU is highly recommended') print('Initializing video data manager') dm = VideoDataManager(use_gpu, **video_dataset_kwargs(args)) trainloader, testloader_dict = dm.return_dataloaders() print('Initializing model: {}'.format(args.arch)) model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, loss={'xent', 'htri'}, pretrained=not args.no_pretrained, use_gpu=use_gpu) print('Model size: {:.3f} M'.format(count_num_param(model))) if args.load_weights and check_isfile(args.load_weights): load_pretrained_weights(model, args.load_weights) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 best_rank1 = checkpoint['rank1'] print('Loaded checkpoint from "{}"'.format(args.resume)) print('- start_epoch: {}\n- rank1: {}'.format(args.start_epoch, best_rank1)) model = nn.DataParallel(model).cuda() if use_gpu else model criterion = CrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth) criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optimizer(model, **optimizer_kwargs(args)) scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs(args)) if args.evaluate: print('Evaluate only') for name in args.target_names: print('Evaluating {} ...'.format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] distmat = test(model, queryloader, galleryloader, args.pool_tracklet_features, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results(distmat, dm.return_testdataset_by_name(name), save_dir=osp.join( args.save_dir, 'ranked_results', name), topk=20) return start_time = time.time() ranklogger = RankLogger(args.source_names, args.target_names) train_time = 0 print('=> Start training') if args.fixbase_epoch > 0: print( 'Train {} for {} epochs while keeping other layers frozen'.format( args.open_layers, args.fixbase_epoch)) initial_optim_state = optimizer.state_dict() for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu, fixbase=True) train_time += round(time.time() - start_train_time) print('Done. All layers are open to train for {} epochs'.format( args.max_epoch)) optimizer.load_state_dict(initial_optim_state) for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_freq > 0 and ( epoch + 1) % args.eval_freq == 0 or (epoch + 1) == args.max_epoch: print('=> Test') for name in args.target_names: print('Evaluating {} ...'.format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] rank1 = test(model, queryloader, galleryloader, args.pool_tracklet_features, use_gpu) ranklogger.write(name, epoch + 1, rank1) save_checkpoint( { 'state_dict': model.state_dict(), 'rank1': rank1, 'epoch': epoch, }, False, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( 'Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.'. format(elapsed, train_time)) ranklogger.show_summary()
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = DatasetManager(dataset_dir=args.dataset, root=args.root) transform_train = T.Compose_Keypt([ T.Random2DTranslation_Keypt((args.width, args.height)), T.RandomHorizontalFlip_Keypt(), T.ToTensor_Keypt(), T.Normalize_Keypt(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose_Keypt([ T.Resize_Keypt((args.width, args.height)), T.ToTensor_Keypt(), T.Normalize_Keypt(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset(dataset.train, keyptaware=args.keyptaware, heatmapaware=args.heatmapaware, segmentaware=args.segmentaware, transform=transform_train, imagesize=(args.width, args.height)), sampler=RandomIdentitySampler(dataset.train, args.train_batch, args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( ImageDataset(dataset.query, keyptaware=args.keyptaware, heatmapaware=args.heatmapaware, segmentaware=args.segmentaware, transform=transform_test, imagesize=(args.width, args.height)), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( ImageDataset(dataset.gallery, keyptaware=args.keyptaware, heatmapaware=args.heatmapaware, segmentaware=args.segmentaware, transform=transform_test, imagesize=(args.width, args.height)), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_vids=dataset.num_train_vids, num_vcolors=dataset.num_train_vcolors, num_vtypes=dataset.num_train_vtypes, keyptaware=args.keyptaware, heatmapaware=args.heatmapaware, segmentaware=args.segmentaware, multitask=args.multitask) print("Model size: {:.3f} M".format(count_num_param(model))) if args.label_smooth: criterion_xent_vid = CrossEntropyLabelSmooth( num_classes=dataset.num_train_vids, use_gpu=use_gpu) criterion_xent_vcolor = CrossEntropyLabelSmooth( num_classes=dataset.num_train_vcolors, use_gpu=use_gpu) criterion_xent_vtype = CrossEntropyLabelSmooth( num_classes=dataset.num_train_vtypes, use_gpu=use_gpu) else: criterion_xent_vid = nn.CrossEntropyLoss() criterion_xent_vcolor = nn.CrossEntropyLoss() criterion_xent_vtype = nn.CrossEntropyLoss() criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.load_weights: # load pretrained weights but ignore layers that don't match in size if check_isfile(args.load_weights): checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format( args.load_weights)) if args.resume: if check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format( args.start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") distmat = test(model, args.keyptaware, args.multitask, queryloader, galleryloader, use_gpu, dataset.vcolor2label, dataset.vtype2label, return_distmat=True) if args.vis_ranked_res: visualize_ranked_results( distmat, dataset, save_dir=osp.join(args.save_dir, 'ranked_results'), topk=100, ) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, args.keyptaware, args.multitask, criterion_xent_vid, criterion_xent_vcolor, criterion_xent_vtype, criterion_htri, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if ((epoch + 1) > args.start_eval and args.eval_step > 0 and (epoch + 1) % args.eval_step == 0) or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, args.keyptaware, args.multitask, queryloader, galleryloader, use_gpu, dataset.vcolor2label, dataset.vtype2label) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() # rank1 = 1 # is_best = True save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.2%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False logger_info = LoggerInfo() sys.stdout = Logger(logger_info) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") # print("Initializing dataset {}".format(args.dataset)) # dataset = data_manager.init_imgreid_dataset( # root=args.root, name=args.dataset, split_id=args.split_id, # cuhk03_labeled=args.cuhk03_labeled, # cuhk03_classic_split=args.cuhk03_classic_split, # ) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False train_dataset = ImageFolder(os.path.join(args.data_dir, "train_all"), transform=transform_train) train_query_dataset = ImageFolder(os.path.join(args.data_dir, "val"), transform=transform_train) train_gallery_dataset = ImageFolder(os.path.join(args.data_dir, "train"), transform=transform_train) test_dataset = ImageFolder(os.path.join(args.data_dir, "probe"), transform=transform_test) query_dataset = ImageFolder(os.path.join(args.data_dir, "query"), transform=transform_test) gallery_dataset = ImageFolder(os.path.join(args.data_dir, "gallery"), transform=transform_test) # train_batch_sampler = VehicleIdBalancedBatchSampler(train_dataset, n_classes=8, n_samples=6) # test_batch_sampler = VehicleIdBalancedBatchSampler(test_dataset, n_classes=8, n_samples=8) train_batch_sampler = VehicleIdCCLBatchSampler(train_dataset, n_classes=n_cls, n_samples=n_samples) trainloader = DataLoader(train_dataset, batch_sampler=train_batch_sampler, num_workers=args.workers, pin_memory=pin_memory) testloader = DataLoader(test_dataset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False) # trainloader = DataLoader( # ImageDataset(dataset.train, transform=transform_train), # batch_sampler=train_batch_sampler, batch_size=args.train_batch, # shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True # ) queryloader = DataLoader( query_dataset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( gallery_dataset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) train_query_loader = DataLoader( train_query_dataset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) train_gallery_loader = DataLoader( train_gallery_dataset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) if args.evaluate: model = models.init_model(name=args.arch, num_classes=len(query_dataset.classes), loss_type=args.loss_type) else: model = models.init_model(name=args.arch, num_classes=len(train_dataset.classes), loss_type=args.loss_type) print("Model size: {:.3f} M".format(count_num_param(model))) if args.label_smooth: criterion = CrossEntropyLabelSmooth(num_classes=len( train_dataset.classes), use_gpu=use_gpu) else: if args.loss_type == 'xent': criterion = nn.CrossEntropyLoss() elif args.loss_type == 'angle': criterion = AngleLoss() elif args.loss_type == 'triplet': # criterion = CoupledClustersLoss(margin=1., triplet_selector=RandomNegativeTripletSelector(margin=1.)) # criterion = OnlineTripletLoss(margin=1., triplet_selector=RandomNegativeTripletSelector(margin=1.)) # criterion = OnlineTripletLoss(margin=1., triplet_selector=HardestNegativeTripletSelector(margin=1.)) criterion = CoupledClustersLoss(margin=1., n_classes=n_cls, n_samples=n_samples) # criterion = OnlineTripletLoss(margin=1., triplet_selector=SemihardNegativeTripletSelector(margin=1.)) elif args.loss_type == 'xent_htri': criterion = XentTripletLoss( margin=1., triplet_selector=RandomNegativeTripletSelector(margin=1.)) else: raise KeyError("Unsupported loss: {}".format(args.loss_type)) # model_param_list = [{'params': model.base.parameters(), 'lr': args.lr}, # {'params': model.classifier.parameters(), 'lr': args.lr * 10}] # optimizer = init_optim(args.optim, model_param_list, lr=1.0, weight_decay=args.weight_decay) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.fixbase_epoch > 0: if hasattr(model, 'classifier') and isinstance(model.classifier, nn.Module): optimizer_tmp = init_optim(args.optim, model.classifier.parameters(), args.fixbase_lr, args.weight_decay) else: print( "Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0" ) args.fixbase_epoch = 0 if args.load_weights: # load pretrained weights but ignore layers that don't match in size if check_isfile(args.load_weights): checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format( args.load_weights)) if args.resume: from functools import partial import pickle pickle.load = partial(pickle.load, encoding="latin1") pickle.Unpickler = partial(pickle.Unpickler, encoding="latin1") if check_isfile(args.resume): checkpoint = torch.load(args.resume) # checkpoint = torch.load(args.resume, map_location=lambda storage, loc: storage, pickle_module=pickle) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format( args.start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() # if args.evaluate: # print("Evaluate only") # distmat = test(model, queryloader, galleryloader, train_query_loader, train_gallery_loader, # use_gpu, return_distmat=True) # if args.vis_ranked_res: # visualize_ranked_results( # distmat, dataset, # save_dir=osp.join(args.save_dir, 'ranked_results'), # topk=20, # ) # return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") if args.fixbase_epoch > 0: print( "Train classifier for {} epochs while keeping base network frozen". format(args.fixbase_epoch)) for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer_tmp, trainloader, use_gpu, freeze_bn=True) train_time += round(time.time() - start_train_time) del optimizer_tmp print("Now open all layers for training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, testloader, queryloader, galleryloader, train_query_loader, train_gallery_loader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, use_gpu_suo=True, fpath=osp.join( args.save_dir, 'checkpoint_ep' + str(epoch + 1) + checkpoint_suffix + '.pth.tar')) print("==> Best Rank-1 {:.2%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): global args, best_rank1 torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_vidreid_dataset(root=args.root, name=args.dataset) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False # decompose tracklets into images for image-based training new_train = [] for img_paths, pid, camid in dataset.train: for img_path in img_paths: new_train.append((img_path, pid, camid)) trainloader = DataLoader( ImageDataset(new_train, transform=transform_train), sampler=RandomIdentitySampler(new_train, args.train_batch, args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( VideoDataset(dataset.query, seq_len=args.seq_len, sample='evenly', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( VideoDataset(dataset.gallery, seq_len=args.seq_len, sample='evenly', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(model))) if args.label_smooth: criterion_xent = CrossEntropyLabelSmooth( num_classes=dataset.num_train_pids, use_gpu=use_gpu) else: criterion_xent = nn.CrossEntropyLoss() criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 best_rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, best_rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") distmat = test(model, queryloader, galleryloader, args.pool, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results( distmat, dataset, save_dir=osp.join(args.save_dir, 'ranked_results'), topk=20, ) return start_time = time.time() train_time = 0 best_epoch = args.start_epoch print("==> Start training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, queryloader, galleryloader, args.pool, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): global args, criterion, testloader_dict, trainloader, use_gpu set_random_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'test.log' if args.evaluate else 'train.log' sys.stdout = Logger(osp.join(args.save_dir, log_name)) print('==========\nArgs:{}\n=========='.format(args)) if use_gpu: print('Currently using GPU {}'.format(args.gpu_devices)) cudnn.benchmark = True else: warnings.warn('Currently using CPU, however, GPU is highly recommended') print('Initializing image data manager') dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args)) trainloader, testloader_dict = dm.return_dataloaders() print('Initializing model: {}'.format(args.arch)) model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, loss={'xent'}, pretrained=not args.no_pretrained, use_gpu=use_gpu) print('Model size: {:.3f} M'.format(count_num_param(model))) if args.load_weights and check_isfile(args.load_weights): load_pretrained_weights(model, args.load_weights) model = nn.DataParallel(model).cuda() if use_gpu else model criterion = CrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth) if args.resume and check_isfile(args.resume): args.start_epoch = resume_from_checkpoint(args.resume, model, optimizer=None) resumed = True else: resumed = False if args.evaluate: print('Evaluate only') for name in args.target_names: print('Evaluating {} ...'.format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results( distmat, dm.return_testdataset_by_name(name), save_dir=osp.join(args.save_dir, 'ranked_results', name), topk=20 ) return time_start = time.time() # ranklogger = RankLogger(args.source_names, args.target_names) print('=> Start training') if not resumed: train_base(model) train_RRI(model, 7) elapsed = round(time.time() - time_start) elapsed = str(datetime.timedelta(seconds=elapsed)) print('Elapsed {}'.format(elapsed))
def main(): global args set_random_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'log_test.txt' if args.evaluate else 'log_train.txt' sys.stdout = Logger(osp.join(args.save_dir, log_name)) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True else: print("Currently using CPU, however, GPU is highly recommended") print("Initializing MultiScale data manager") assert args.train_batch_size % args.train_loss_batch_size == 0, "'{}' is not divisable by {}".format(args.train_loss_batch_size, args.train_loss_batch_size) dm = ImageDataManager(use_gpu, scales=[224,160], **image_dataset_kwargs(args)) trainloader, testloader_dict = dm.return_dataloaders() # sys.exit(0) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, input_size=args.width, loss={'xent'}, use_gpu=use_gpu) print("Model size: {:.3f} M".format(count_num_param(model))) # print(model) criterion = CrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth) optimizer = init_optimizer(model.parameters(), **optimizer_kwargs(args)) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) # # scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=3, verbose=True, threshold=1e-04) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()} model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, checkpoint['rank1'])) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] test_set = dm.return_testdataset_by_name(name) rank1, mAP = test(model, test_set, name, queryloader, galleryloader, use_gpu, visualize=args.visualize_ranks) return start_time = time.time() ranklogger = RankLogger(args.source_names, args.target_names) maplogger = RankLogger(args.source_names, args.target_names) train_time = 0 # Tensorboard writer = SummaryWriter(log_dir=osp.join('runs', args.save_dir)) print("=> Start training") if args.fixbase_epoch > 0: print("Train {} for {} epochs while keeping other layers frozen".format(args.open_layers, args.fixbase_epoch)) initial_optim_state = optimizer.state_dict() for epoch in range(args.fixbase_epoch): start_train_time = time.time() loss, prec1 = train(epoch, model, criterion, optimizer, trainloader, writer, use_gpu, fixbase=True) writer.add_scalar('train/loss', loss, epoch+1) writer.add_scalar('train/prec1', prec1, epoch+1) print('Epoch: [{:02d}] [Average Loss:] {:.4f}\t [Average Prec.:] {:.2%}'.format(epoch+1, loss, prec1)) train_time += round(time.time() - start_train_time) print("Done. All layers are open to train for {} epochs".format(args.max_epoch)) optimizer.load_state_dict(initial_optim_state) args.start_epoch += args.fixbase_epoch args.max_epoch += args.fixbase_epoch for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() loss, prec1 = train(epoch, model, criterion, optimizer, trainloader, writer, use_gpu) writer.add_scalar('train/loss', loss, epoch+1) writer.add_scalar('train/prec1', prec1, epoch+1) print('Epoch: [{:02d}] [Average Loss:] {:.4f}\t [Average Prec.:] {:.2%}'.format(epoch+1, loss, prec1)) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_freq > 0 and (epoch + 1) % args.eval_freq == 0 or (epoch + 1) == args.max_epoch: print("=> Test") for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] test_set = dm.return_testdataset_by_name(name) if epoch+1 == args.max_epoch: rank1, mAP = test(model, test_set, name, queryloader, galleryloader, use_gpu, visualize=True) else: rank1, mAP = test(model, test_set, name, queryloader, galleryloader, use_gpu) writer.add_scalar(name + '_test/top1', rank1, epoch+1) writer.add_scalar(name + '_test/mAP', mAP, epoch+1) ranklogger.write(name, epoch + 1, rank1) maplogger.write(name, epoch + 1, mAP) if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, False, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) # save last checkpoint save_checkpoint({ 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, False, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time)) ranklogger.show_summary() maplogger.show_summary()
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset(root=args.root, name=args.dataset, split_id=args.split_id) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) testloader = DataLoader( ImageDataset(dataset.test, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, loss={'xent'}, use_gpu=use_gpu) print("Model size: {:.3f} M".format(count_num_param(model))) gender_criterion_xent = nn.CrossEntropyLoss() staff_criterion_xent = nn.CrossEntropyLoss() customer_criterion_xent = nn.CrossEntropyLoss() stand_criterion_xent = nn.CrossEntropyLoss() sit_criterion_xent = nn.CrossEntropyLoss() phone_criterion_xent = nn.CrossEntropyLoss() optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if use_gpu: model = nn.DataParallel(model).cuda() start_time = time.time() train_time = 0 best_score = 0 best_epoch = args.start_epoch print("==> Start training") ################################### 修改到这里,把train 和 test改一下就好 for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, gender_criterion_xent, staff_criterion_xent, customer_criterion_xent, \ stand_criterion_xent, sit_criterion_xent, phone_criterion_xent, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") gender_accurary, staff_accurary, customer_accurary, stand_accurary, sit_accurary, phone_accurary = test( model, testloader, use_gpu) Score = (gender_accurary + staff_accurary + customer_accurary + stand_accurary + sit_accurary + phone_accurary) * 100 is_best = Score > best_score if is_best: best_score = Score best_gender_acc = gender_accurary best_staff_acc = staff_accurary best_customer_acc = customer_accurary best_stand_acc = stand_accurary best_sit_acc = sit_accurary best_phone_acc = phone_accurary best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': Score, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print( "==> Best best_score {} |Gender_acc {}\t Staff_acc {}\t Customer_acc {}\t Stand_acc {}\t Sit_acc {}\t Phone_acc {}|achieved at epoch {}" .format(best_score, best_gender_acc, best_staff_acc, best_customer_acc, best_stand_acc, best_sit_acc, best_phone_acc, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): global args, dropout_optimizer torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'log_test.txt' if args.evaluate else 'log_train.txt' sys.stderr = sys.stdout = Logger(osp.join(args.save_dir, log_name)) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU, however, GPU is highly recommended") print("Initializing image data manager") dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args)) trainloader, testloader_dict = dm.return_dataloaders() print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, loss={'xent'}, use_gpu=use_gpu, dropout_optimizer=dropout_optimizer) print(model) print("Model size: {:.3f} M".format(count_num_param(model))) # criterion = WrappedCrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth) criterion, fix_criterion, switch_criterion, htri_param_controller = get_criterions(dm.num_train_pids, use_gpu, args) regularizer, reg_param_controller = get_regularizer(args.regularizer) optimizer = init_optimizer(model.parameters(), **optimizer_kwargs(args)) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size try: checkpoint = torch.load(args.load_weights) except Exception as e: print(e) checkpoint = torch.load(args.load_weights, map_location={'cuda:0': 'cpu'}) # dropout_optimizer.set_p(checkpoint.get('dropout_p', 0)) # print(list(checkpoint.keys()), checkpoint['dropout_p']) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()} model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) state = model.state_dict() state.update(checkpoint['state_dict']) model.load_state_dict(state) # args.start_epoch = checkpoint['epoch'] + 1 print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, checkpoint['rank1'])) if use_gpu: model = nn.DataParallel(model, device_ids=list(range(len(args.gpu_devices.split(','))))).cuda() extract_train_info(model, trainloader)
def main(): global args if not args.evaluate: raise RuntimeError('Test only!') torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'log_test.txt' if args.evaluate else 'log_train.txt' log_fn = osp.join(args.save_dir, log_name) sys.stderr = sys.stdout = Logger(log_fn) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU, however, GPU is highly recommended") print("Initializing image data manager") dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args)) trainloader, testloader_dict = dm.return_dataloaders() print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, loss={'xent'}, use_gpu=use_gpu, args=vars(args)) print(model) print("Model size: {:.3f} M".format(count_num_param(model))) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size try: checkpoint = torch.load(args.load_weights) except Exception as e: print(e) checkpoint = torch.load(args.load_weights, map_location={'cuda:0': 'cpu'}) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) state = model.state_dict() state.update(checkpoint['state_dict']) model.load_state_dict(state) # args.start_epoch = checkpoint['epoch'] + 1 print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, checkpoint['rank1'])) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") for name in args.target_names: print("Evaluating {} ...".format(name)) distmat = test(model, testloader_dict[name], use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results(distmat, dm.return_testdataset_by_name(name), save_dir=osp.join( args.save_dir, 'ranked_results', name), topk=20) return
def main(): global args, dropout_optimizer torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'log_test.txt' if args.evaluate else 'log_train.txt' sys.stderr = sys.stdout = Logger(osp.join(args.save_dir, log_name)) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU, however, GPU is highly recommended") print("Initializing image data manager") dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args)) trainloader, testloader_dict = dm.return_dataloaders() print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, loss={'xent'}, use_gpu=use_gpu, dropout_optimizer=dropout_optimizer, args=vars(args)) print(model) print("Model size: {:.3f} M".format(count_num_param(model))) # criterion = WrappedCrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth) criterion, fix_criterion, switch_criterion, htri_param_controller = get_criterions( dm.num_train_pids, use_gpu, args) regularizer, reg_param_controller = get_regularizer(args.regularizer) optimizer = init_optimizer(model.parameters(), **optimizer_kwargs(args)) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size try: checkpoint = torch.load(args.load_weights) except Exception as e: print(e) checkpoint = torch.load(args.load_weights, map_location={'cuda:0': 'cpu'}) # dropout_optimizer.set_p(checkpoint.get('dropout_p', 0)) # print(list(checkpoint.keys()), checkpoint['dropout_p']) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) state = model.state_dict() state.update(checkpoint['state_dict']) model.load_state_dict(state) # args.start_epoch = checkpoint['epoch'] + 1 print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, checkpoint['rank1'])) if use_gpu: model = nn.DataParallel( model, device_ids=list(range(len(args.gpu_devices.split(','))))).cuda() if args.evaluate: print("Evaluate only") for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'], testloader_dict[ name]['query_flip'] galleryloader = testloader_dict[name]['gallery'], testloader_dict[ name]['gallery_flip'] distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results(distmat, dm.return_testdataset_by_name(name), save_dir=osp.join( args.save_dir, 'ranked_results', name), topk=20) return start_time = time.time() ranklogger = RankLogger(args.source_names, args.target_names) train_time = 0 print("==> Start training") if os.environ.get('test_first') is not None: for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'], testloader_dict[ name]['query_flip'] galleryloader = testloader_dict[name]['gallery'], testloader_dict[ name]['gallery_flip'] rank1 = test(model, queryloader, galleryloader, use_gpu) if args.fixbase_epoch > 0: oldenv = os.environ.get('sa', '') os.environ['sa'] = '' print( "Train {} for {} epochs while keeping other layers frozen".format( args.open_layers, args.fixbase_epoch)) initial_optim_state = optimizer.state_dict() for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, fix_criterion, regularizer, optimizer, trainloader, use_gpu, fixbase=True) train_time += round(time.time() - start_train_time) print("Done. All layers are open to train for {} epochs".format( args.max_epoch)) optimizer.load_state_dict(initial_optim_state) os.environ['sa'] = oldenv max_r1 = 0 for epoch in range(args.start_epoch, args.max_epoch): dropout_optimizer.set_epoch(epoch) reg_param_controller.set_epoch(epoch) htri_param_controller.set_epoch(epoch) dropout_optimizer.set_training(True) start_train_time = time.time() print(epoch, args.switch_loss) print(criterion) cond = args.switch_loss > 0 and epoch >= args.switch_loss cond = cond or (args.switch_loss < 0 and args.switch_loss + args.max_epoch < epoch) if cond: print('Switch!') criterion = switch_criterion train(epoch, model, criterion, regularizer, optimizer, trainloader, use_gpu, fixbase=False, switch_loss=cond) train_time += round(time.time() - start_train_time) if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': 0, 'epoch': epoch, 'dropout_p': dropout_optimizer.p, }, False, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_freq > 0 and ( epoch + 1) % args.eval_freq == 0 or (epoch + 1) == args.max_epoch: print("==> Test") dropout_optimizer.set_training(False) # IMPORTANT! for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'], testloader_dict[ name]['query_flip'] galleryloader = testloader_dict[name][ 'gallery'], testloader_dict[name]['gallery_flip'] print('!!!!!!!!FC!!!!!!!!') os.environ['NOFC'] = '' rank1 = test(model, queryloader, galleryloader, use_gpu) ranklogger.write(name, epoch + 1, rank1) if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() if max_r1 < rank1: print('Save!', max_r1, rank1) save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, 'dropout_p': dropout_optimizer.p, }, False, osp.join(args.save_dir, 'checkpoint_best.pth.tar')) max_r1 = rank1 elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time)) ranklogger.show_summary()
def main(): global args torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'log_test.txt' if args.evaluate else 'log_train.txt' sys.stdout = Logger(osp.join(args.save_dir, log_name)) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU, however, GPU is highly recommended") print("Initializing image data manager") dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args)) trainloader, testloader_dict = dm.return_dataloaders() # ReID-Stream: print("Initializing ReID-Stream: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, reid_dim=args.reid_dim, loss={'xent', 'htri'}) print("ReID Model size: {:.3f} M".format(count_num_param(model))) criterion_xent = CrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth) criterion_htri = TripletLoss(margin=args.margin) # 2. Optimizer # Main ReID-Stream: optimizer = init_optimizer(model.parameters(), **optimizer_kwargs(args)) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results(distmat, dm.return_testdataset_by_name(name), save_dir=osp.join( args.save_dir, 'ranked_results', name), topk=20) return start_time = time.time() ranklogger = RankLogger(args.source_names, args.target_names) train_time = 0 print("==> Start training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, \ optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_freq > 0 and ( epoch + 1) % args.eval_freq == 0 or (epoch + 1) == args.max_epoch: print("==> Test") for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] rank1 = test(model, queryloader, galleryloader, use_gpu) ranklogger.write(name, epoch + 1, rank1) if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, False, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time)) ranklogger.show_summary()
type=int, default=128, help="width of an image (default: 128)") parser.add_argument( '--load-weights', type=str, default='', help="load pretrained weights but ignores layers that don't match in size") args = parser.parse_args() print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=10, loss={'htri'}, use_gpu=True) print("Model size: {:.3f} M".format(count_num_param(model))) checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) model.eval() dummy_input = torch.randn(100, 3, args.height, args.width, requires_grad=True)
def main(): global args torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'log_test.txt' sys.stdout = Logger(osp.join(args.save_dir, log_name)) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU, however, GPU is highly recommended") print("Initializing image data manager") if not args.convert_to_onnx: # and not args.infer: dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args)) trainloader, trainloader_dict, testloader_dict = dm.return_dataloaders( ) num_train_pids = 100 print("Initializing model: {}".format(args.arch)) model = models.init_model( name=args.arch, num_classes=num_train_pids, loss={'xent', 'htri'}, pretrained=False if args.load_weights else 'imagenet', grayscale=args.grayscale, ceil_mode=not args.convert_to_onnx, infer=True, bits=args.bits, normalize_embeddings=args.normalize_embeddings, normalize_fc=args.normalize_fc, convbn=args.convbn) print("Model size: {:.3f} M".format(count_num_param(model))) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size load_weights(model, args.load_weights) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.absorb_bn: search_absorbed_bn(model) if args.quantization or args.save_quantized_model: from gap_quantization.quantization import ModelQuantizer from gap_quantization.dump_utils import dump_quant_params, remove_extra_dump, remove_cat_files if args.quant_data_dir is None: raise AttributeError('quant-data-dir argument is required.') num_channels = 1 if args.grayscale else 3 cfg = { "bits": args.bits, # number of bits to store weights and activations "accum_bits": 32, # number of bits to store intermediate convolution result "signed": True, # use signed numbers "save_folder": args.save_dir, # folder to save results "data_source": args. quant_data_dir, # folder with images to collect dataset statistics "use_gpu": False, # use GPU for inference "batch_size": 1, "num_workers": 0, # number of workers for PyTorch dataloader "verbose": True, "save_params": args. save_quantized_model, # save quantization parameters to the file "quantize_forward": True, # replace usual convs, poolings, ... with GAP-like ones "num_input_channels": num_channels, "raw_input": args.no_normalize, "double_precision": args.double_precision # use double precision convolutions } model = model.cpu() quantizer = ModelQuantizer( model, cfg, dm.transform_test ) # transform test is OK if we use args.no_normalize quantizer.quantize_model( ) # otherwise we need to add QuantizeInput operation if args.infer: if args.image_path == '': raise AttributeError('Image for inference is required') quantizer.dump_activations(args.image_path, dm.transform_test, save_dir=os.path.join( args.save_dir, 'activations_dump')) dump_quant_params(args.save_dir, args.convbn) if args.convbn: remove_extra_dump( os.path.join(args.save_dir, 'activations_dump')) remove_cat_files(args.save_dir) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") for name in args.target_names: if not 'lfw' in name.lower(): print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] distmat = test(args, model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results( distmat, dm.return_testdataset_by_name(name), save_dir=osp.join(args.save_dir, 'ranked_results', name), topk=20) else: model.eval() same_acc, diff_acc, all_acc, auc, thresh = evaluate( args, dm.lfw_dataset, model, compute_embeddings_lfw, args.test_batch_size, verbose=False, show_failed=args.show_failed, load_embeddings=args.load_embeddings) log.info('Validation accuracy: {0:.4f}, {1:.4f}'.format( same_acc, diff_acc)) log.info('Validation accuracy mean: {0:.4f}'.format(all_acc)) log.info('Validation AUC: {0:.4f}'.format(auc)) log.info('Estimated threshold: {0:.4f}'.format(thresh)) #roc_auc(model, '/home/maxim/data/lfw/pairsTest.txt', '/media/slow_drive/cropped_lfw', args, use_gpu) return
def main(): global args torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'log_test.txt' if args.evaluate else 'log_train.txt' sys.stdout = Logger(osp.join(args.save_dir, log_name)) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU, however, GPU is highly recommended") print("Initializing image data manager") dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args)) trainloader, testloader_dict = dm.return_dataloaders() print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, loss={'xent', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(model))) criterion_xent = CrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth) criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optimizer(model.parameters(), **optimizer_kwargs(args)) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()} model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, checkpoint['rank1'])) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results( distmat, dm.return_testdataset_by_name(name), save_dir=osp.join(args.save_dir, 'ranked_results', name), topk=20 ) return start_time = time.time() ranklogger = RankLogger(args.source_names, args.target_names) train_time = 0 print("=> Start training") if args.fixbase_epoch > 0: print("Train {} for {} epochs while keeping other layers frozen".format(args.open_layers, args.fixbase_epoch)) initial_optim_state = optimizer.state_dict() for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu, fixbase=True) train_time += round(time.time() - start_train_time) print("Done. All layers are open to train for {} epochs".format(args.max_epoch)) optimizer.load_state_dict(initial_optim_state) for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_freq > 0 and (epoch + 1) % args.eval_freq == 0 or (epoch + 1) == args.max_epoch: print("=> Test") for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'] galleryloader = testloader_dict[name]['gallery'] rank1 = test(model, queryloader, galleryloader, use_gpu) ranklogger.write(name, epoch + 1, rank1) if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, False, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time)) ranklogger.show_summary()
def main(args): args = parser.parse_args(args) #global best_rank1 best_rank1 = -np.inf torch.manual_seed(args.seed) # np.random.seed(args.seed) # random.seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: test_dir = args.save_dir if args.save_dir == 'log': if args.resume: test_dir = os.path.dirname(args.resume) else: test_dir = os.path.dirname(args.load_weights) sys.stdout = Logger(osp.join(test_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset( root=args.root, name=args.dataset, split_id=args.split_id, cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split, split_wild=args.split_wild) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), #T.Resize((args.height, args.width)), T.RandomSizedEarser(), T.RandomHorizontalFlip_custom(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( ImageDataset(dataset.query, transform=transform_test, return_path=args.draw_tsne), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( ImageDataset(dataset.gallery, transform=transform_test, return_path=args.draw_tsne), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model( name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'angular'} if args.use_angular else {'xent'}, use_gpu=use_gpu) print("Model size: {:.3f} M".format(count_num_param(model))) use_autoTune = False if not (args.use_angular): if args.label_smooth: print("Using Label Smoothing with epsilon", args.label_epsilon) criterion = CrossEntropyLabelSmooth( num_classes=dataset.num_train_pids, epsilon=args.label_epsilon, use_gpu=use_gpu) elif args.focal_loss: print("Using Focal Loss with gamma=", args.focal_gamma) criterion = FocalLoss(gamma=args.focal_gamma) else: print("Using Normal Cross-Entropy") criterion = nn.CrossEntropyLoss() if args.jsd: print("Using JSD regularizer") criterion = (criterion, JSD_loss(dataset.num_train_pids)) if args.auto_tune_mtl: print("Using AutoTune") use_autoTune = True criterion = MultiHeadLossAutoTune( list(criterion), [args.lambda_xent, args.confidence_beta]).cuda() else: if args.confidence_penalty: print("Using Confidence Penalty", args.confidence_beta) criterion = (criterion, ConfidencePenalty()) if args.auto_tune_mtl and args.confidence_penalty: print("Using AutoTune") use_autoTune = True criterion = MultiHeadLossAutoTune( list(criterion), [args.lambda_xent, -args.confidence_beta]).cuda() else: if args.label_smooth: print("Using Angular Label Smoothing") criterion = AngularLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) else: print("Using Angular Loss") criterion = AngleLoss() if use_autoTune: optimizer = init_optim( args.optim, list(model.parameters()) + list(criterion.parameters()), args.lr, args.weight_decay) else: optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) if args.scheduler: scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.fixbase_epoch > 0: if hasattr(model, 'classifier') and isinstance(model.classifier, nn.Module): if use_autoTune: optimizer_tmp = init_optim( args.optim, list(model.classifier.parameters()) + list(criterion.parameters()), args.fixbase_lr, args.weight_decay) else: optimizer_tmp = init_optim(args.optim, model.classifier.parameters(), args.fixbase_lr, args.weight_decay) else: print( "Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0" ) args.fixbase_epoch = 0 if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 best_rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, best_rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.single_folder != '': extract_features(model, use_gpu, args, transform_test, return_distmat=False) return if args.evaluate: print("Evaluate only") test_dir = args.save_dir if args.save_dir == 'log': if args.resume: test_dir = os.path.dirname(args.resume) else: test_dir = os.path.dirname(args.load_weights) distmat = test(model, queryloader, galleryloader, use_gpu, args, writer=None, epoch=-1, return_distmat=True, tsne_clusters=args.tsne_labels) if args.visualize_ranks: visualize_ranked_results( distmat, dataset, save_dir=osp.join(test_dir, 'ranked_results'), topk=10, ) return writer = SummaryWriter(log_dir=osp.join(args.save_dir, 'tensorboard')) start_time = time.time() train_time = 0 best_epoch = args.start_epoch print("==> Start training") if args.fixbase_epoch > 0: print( "Train classifier for {} epochs while keeping base network frozen". format(args.fixbase_epoch)) for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer_tmp, trainloader, use_gpu, writer, args, freeze_bn=True) train_time += round(time.time() - start_train_time) del optimizer_tmp print("Now open all layers for training") best_epoch = 0 for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer, trainloader, use_gpu, writer, args) train_time += round(time.time() - start_train_time) if args.scheduler: scheduler.step() if (epoch + 1) > args.start_eval and ( (args.save_epoch > 0 and (epoch + 1) % args.save_epoch == 0) or (args.eval_step > 0 and (epoch + 1) % args.eval_step == 0) or (epoch + 1) == args.max_epoch): if (epoch + 1) == args.max_epoch: if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': -1, 'epoch': epoch, }, False, osp.join( args.save_dir, 'beforeTesting_checkpoint_ep' + str(epoch + 1) + '.pth.tar')) is_best = False rank1 = -1 if args.eval_step > 0: print("==> Test") rank1 = test(model, queryloader, galleryloader, use_gpu, args, writer=writer, epoch=epoch) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time)) return best_rank1, best_epoch
def main(): global use_apex global args torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False log_name = 'log_test.txt' if args.evaluate else 'log_train.txt' sys.stderr = sys.stdout = Logger(osp.join(args.save_dir, log_name)) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU, however, GPU is highly recommended") print("Initializing image data manager") dm = ImageDataManager(use_gpu, **image_dataset_kwargs(args)) trainloader, testloader_dict = dm.return_dataloaders() print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, loss={'xent'}, use_gpu=use_gpu, args=vars(args)) print(model) print("Model size: {:.3f} M".format(count_num_param(model))) if use_gpu: print("using gpu") model = model.cuda() print("criterion===>") criterion = get_criterion(dm.num_train_pids, use_gpu, args) print(criterion) print("regularizer===>") regularizer = get_regularizer(vars(args)) print(regularizer) print("optimizer===>") optimizer = init_optimizer(model.parameters(), **optimizer_kwargs(args)) print(optimizer) print("scheduler===>") scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'max', factor=0.1, patience=5, verbose=True) print(scheduler) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size try: checkpoint = torch.load(args.load_weights) except Exception as e: print(e) checkpoint = torch.load(args.load_weights, map_location={'cuda:0': 'cpu'}) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) max_r1 = 0 if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) state = model.state_dict() state.update(checkpoint['state_dict']) model.load_state_dict(state) optimizer.load_state_dict(checkpoint['optimizer']) args.start_epoch = checkpoint['epoch'] + 1 max_r1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, checkpoint['rank1'])) if use_apex: print("using apex") model, optimizer = amp.initialize(model, optimizer, opt_level="O0") if args.evaluate: print("Evaluate only") for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'], testloader_dict[ name]['query_flip'] galleryloader = testloader_dict[name]['gallery'], testloader_dict[ name]['gallery_flip'] distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results(distmat, dm.return_testdataset_by_name(name), save_dir=osp.join( args.save_dir, 'ranked_results', name), topk=20) return start_time = time.time() ranklogger = RankLogger(args.source_names, args.target_names) train_time = 0 print("==> Start training") if args.fixbase_epoch > 0: oldenv = os.environ.get('sa', '') os.environ['sa'] = '' print( "Train {} for {} epochs while keeping other layers frozen".format( args.open_layers, args.fixbase_epoch)) initial_optim_state = optimizer.state_dict() for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion, regularizer, optimizer, trainloader, use_gpu, fixbase=True) train_time += round(time.time() - start_train_time) print("Done. All layers are open to train for {} epochs".format( args.max_epoch)) optimizer.load_state_dict(initial_optim_state) os.environ['sa'] = oldenv for epoch in range(args.start_epoch, args.max_epoch): auto_reset_learning_rate(optimizer, args) print( f"===========================start epoch {epoch + 1} {now()}===========================================" ) print(f"lr:{optimizer.param_groups[0]['lr']}") loss = train(epoch, model, criterion, regularizer, optimizer, trainloader, use_gpu, fixbase=False) train_time += round(time.time() - start_train_time) state_dict = model.state_dict() rank1 = 0 if (epoch + 1) > args.start_eval and args.eval_freq > 0 and ( epoch + 1) % args.eval_freq == 0 or (epoch + 1) == args.max_epoch: print("==> Test") for name in args.target_names: print("Evaluating {} ...".format(name)) queryloader = testloader_dict[name]['query'], testloader_dict[ name]['query_flip'] galleryloader = testloader_dict[name][ 'gallery'], testloader_dict[name]['gallery_flip'] rank1 = test(model, queryloader, galleryloader, use_gpu) ranklogger.write(name, epoch + 1, rank1) if max_r1 < rank1: print('Save!', max_r1, rank1) save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, 'optimizer': optimizer.state_dict(), }, False, osp.join(args.save_dir, 'checkpoint_best.pth.tar')) max_r1 = rank1 save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, 'optimizer': optimizer.state_dict(), }, False, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) scheduler.step(rank1) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time)) ranklogger.show_summary()
def main(args): args = parser.parse_args(args) #global best_rank1 best_rank1 = -np.inf torch.manual_seed(args.seed) # np.random.seed(args.seed) # random.seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: test_dir = args.save_dir if args.save_dir == 'log': if args.resume: test_dir = os.path.dirname(args.resume) else: test_dir = os.path.dirname(args.load_weights) sys.stdout = Logger(osp.join(test_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) # print("Currently using GPU {}".format(args.gpu_devices)) # #cudnn.benchmark = False # cudnn.deterministic = True # torch.cuda.manual_seed_all(args.seed) # torch.set_default_tensor_type('torch.DoubleTensor') else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset( root=args.root, name=args.dataset, split_id=args.split_id, cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split, ) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), #T.Resize((args.height, args.width)), #T.RandomSizedEarser(), T.RandomHorizontalFlip(), #T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False if 'stanford' in args.dataset: datasetLoader = ImageDataset_stanford else: datasetLoader = ImageDataset if args.crop_img: print("Using Cropped Images") else: print("NOT using cropped Images") trainloader = DataLoader( datasetLoader(dataset.train, -1, crop=args.crop_img, transform=transform_train), batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) testloader = DataLoader( datasetLoader(dataset.test, -1, crop=args.crop_img, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model( name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'angular'} if args.use_angular else {'xent'}, use_gpu=use_gpu) print("Model size: {:.3f} M".format(count_num_param(model))) if not (args.use_angular): if args.label_smooth: print("Using Label Smoothing") criterion = CrossEntropyLabelSmooth( num_classes=dataset.num_train_pids, use_gpu=use_gpu) else: criterion = nn.CrossEntropyLoss() else: if args.label_smooth: print("Using Label Smoothing") criterion = AngularLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) else: criterion = AngleLoss() optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) if args.scheduler != 0: scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.fixbase_epoch > 0: if hasattr(model, 'classifier') and isinstance(model.classifier, nn.Module): optimizer_tmp = init_optim( args.optim, list(model.classifier.parameters()) + list(model.encoder.parameters()), args.fixbase_lr, args.weight_decay) else: print( "Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0" ) args.fixbase_epoch = 0 if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 best_rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, best_rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test_dir = args.save_dir if args.save_dir == 'log': if args.resume: test_dir = os.path.dirname(args.resume) else: test_dir = os.path.dirname(args.load_weights) distmat = test(model, testloader, use_gpu, args, writer=None, epoch=-1, return_distmat=True, draw_tsne=args.draw_tsne, tsne_clusters=args.tsne_labels, use_cosine=args.plot_deltaTheta) if args.visualize_ranks: visualize_ranked_results( distmat, dataset, save_dir=osp.join(test_dir, 'ranked_results'), topk=10, ) if args.plot_deltaTheta: plot_deltaTheta(distmat, dataset, save_dir=osp.join(test_dir, 'deltaTheta_results'), min_rank=1) return writer = SummaryWriter(log_dir=osp.join(args.save_dir, 'tensorboard')) start_time = time.time() train_time = 0 best_epoch = args.start_epoch print("==> Start training") if args.test_rot: print("Training only classifier for rotation") model = models.init_model(name='rot_tester', base_model=model, inplanes=2048, num_rot_classes=8) criterion_rot = nn.CrossEntropyLoss() optimizer_rot = init_optim(args.optim, model.fc_rot.parameters(), args.fixbase_lr, args.weight_decay) if use_gpu: model = nn.DataParallel(model).cuda() try: best_epoch = 0 for epoch in range(0, args.max_epoch): start_train_time = time.time() train_rotTester(epoch, model, criterion_rot, optimizer_rot, trainloader, use_gpu, writer, args) train_time += round(time.time() - start_train_time) if args.scheduler != 0: scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or ( epoch + 1) == args.max_epoch: if (epoch + 1) == args.max_epoch: if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': -1, 'epoch': epoch, }, False, osp.join( args.save_dir, 'beforeTesting_checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Test") rank1 = test_rotTester(model, criterion_rot, queryloader, galleryloader, trainloader, use_gpu, args, writer=writer, epoch=epoch) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Cccuracy {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}." .format(elapsed, train_time)) return best_rank1, best_epoch except KeyboardInterrupt: if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': -1, 'epoch': epoch, }, False, osp.join( args.save_dir, 'keyboardInterrupt_checkpoint_ep' + str(epoch + 1) + '.pth.tar')) return None, None if args.fixbase_epoch > 0: print( "Train classifier for {} epochs while keeping base network frozen". format(args.fixbase_epoch)) for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer_tmp, trainloader, use_gpu, writer, args, freeze_bn=True) train_time += round(time.time() - start_train_time) del optimizer_tmp print("Now open all layers for training") best_epoch = 0 for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer, trainloader, use_gpu, writer, args) train_time += round(time.time() - start_train_time) if args.scheduler != 0: scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: if (epoch + 1) == args.max_epoch: if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': -1, 'epoch': epoch, }, False, osp.join( args.save_dir, 'beforeTesting_checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Test") rank1 = test(model, testloader, use_gpu, args, writer=writer, epoch=epoch) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time)) return best_rank1, best_epoch
def main(): torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset( root=args.root, name=args.dataset, split_id=args.split_id, cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split, ) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( ImageDataset(dataset.query, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( ImageDataset(dataset.gallery, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent'}, use_gpu=use_gpu) print("Model size: {:.3f} M".format(count_num_param(model))) if args.label_smooth: criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) else: criterion = nn.CrossEntropyLoss() optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.fixbase_epoch > 0: if hasattr(model, 'classifier') and isinstance(model.classifier, nn.Module): optimizer_tmp = init_optim(args.optim, model.classifier.parameters(), args.fixbase_lr, args.weight_decay) else: print( "Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0" ) args.fixbase_epoch = 0 if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 best_rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, best_rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.vis_ranked_res: visualize_ranked_results( distmat, dataset, save_dir=osp.join(args.save_dir, 'ranked_results'), topk=20, ) return start_time = time.time() train_time = 0 best_epoch = args.start_epoch print("==> Start training") if args.fixbase_epoch > 0: print( "Train classifier for {} epochs while keeping base network frozen". format(args.fixbase_epoch)) for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer_tmp, trainloader, use_gpu, freeze_bn=True) train_time += round(time.time() - start_train_time) del optimizer_tmp print("Now open all layers for training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, queryloader, galleryloader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): ################################# Setting gpu and out file ################################################################### if args.use_cpu: use_gpu = False use_gpu = 1 if use_gpu == 1: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices os.environ['PYTHONHASHSEED'] = str(args.seed) if args.is_deterministic == 1: print("Currently using GPU {}".format(args.gpu_devices)) print('is deterministic') np.random.seed(args.seed) torch.manual_seed(args.seed) random.seed(args.seed) torch.cuda.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False use_gpu = torch.cuda.is_available() else: print("Currently using CPU (GPU is highly recommended)") if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) print("Initializing dataset {}".format(args.dataset)) [dataset, trainloader_reid, queryloader, galleryloader] = initialize_loader(use_gpu) ################################### Initialize model ############################################################### if args.num_classes_attributes == 0: num_classes_attributes = (6, 4, 11, 14, 6) if args.num_classes_attributes > 1: num_classes_attributes = (33, 6, 58, 65, 9) if args.dataset == 'pa100K': num_classes_attributes = (28, 9, 3, 14, 24, 9) if args.num_classes_attributes == 0: num_classes_attributes = (5, 3, 3, 6, 9, 10) if args.dataset == 'dukemtmcreid': num_classes_attributes = (15, 2, 47, 62, 9) if args.num_classes_attributes == 0: num_classes_attributes = (2, 2, 8, 11, 6) num_group_attributes = len(num_classes_attributes) print("Initializing model: {}".format(args.arch)) if args.arch == 'resnetAttW2VText': dim_features = 50 * len(num_classes_attributes) model = models.init_model(args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}, num_group_attributes=num_group_attributes, dim_features=50, glove=True) criterion_attributes = LossW2V( num_classes_attributes=num_classes_attributes, attr_loss_type=args.attr_loss_type) if args.arch == 'resnetAttW2VAttributes': dim_features = sum(num_classes_attributes) model = models.init_model( args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}, num_group_attributes=num_group_attributes, num_classes_attributes=num_classes_attributes, dim_features=dim_features, glove=True) criterion_attributes = LossW2V( num_classes_attributes=num_classes_attributes) all_parameters = model.parameters() optimizer = init_optim(args.optim, all_parameters, args.lr, args.weight_decay) ################################### Loss functions ############################## print("Model size: {:.3f} M".format(count_num_param(model))) criterion_htri_reid = TripletLoss(margin=0.3) criterion_xent_reid = nn.CrossEntropyLoss() ################################### Pretrained models ############################## if args.load_weights: load_weights(model) if args.resume: resume(model) if use_gpu: model = nn.DataParallel(model).cuda() ################################### Only evaluation ############################### writer = SummaryWriter(log_dir=osp.join(args.exp_dir, 'tensorboard')) if args.evaluate: print("Evaluate by test") test_and_evaluate_dist_mat(writer, model, queryloader, galleryloader, use_gpu=use_gpu, save_features=True, load_features=False, arch=args.arch, size=dim_features) return ################################### Training ###################################### start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 # Schedule learning rate print("==> Start training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() if args.is_warmup: adjust_lr(epoch, optimizer, args.lr) if args.arch == 'resnetAttW2VText': train_w2v_single_batch_text(writer, epoch, model, criterion_htri_reid, criterion_attributes, optimizer, trainloader_reid, use_gpu) if args.arch == 'resnetAttW2VAttributes': train_w2v_single_batch_attributes(writer, epoch, model, criterion_htri_reid, criterion_attributes, optimizer, trainloader_reid, use_gpu) train_time += round(time.time() - start_train_time) if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time)) state_dict = model.module.state_dict() save_checkpoint(state={ 'state_dict': state_dict, 'epoch': epoch, }, fpath=osp.join( args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("Evaluate by test") test_and_evaluate_dist_mat(writer, model, queryloader, galleryloader, use_gpu=use_gpu, save_features=True, load_features=False, arch=args.arch, size=dim_features)