def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir + '/log')) if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else (256, 128) dataset, num_classes, train_loader, trainvallabel, val_loader, query_loader, gallery_loader, query_loader_s, gallery_loader_s = \ get_data(args.dataset, args.split, args.data_dir, args.height, args.width, args.batch_size, args.num_instances, args.workers, args.combine_trainval, args.flip_prob, args.padding, args.re_prob) print(num_classes) gen_a = AdaINGen(3, num_classes) gen_b = AdaINGen(3, num_classes) # id_a = ft_netAB(num_classes, stride=1, norm="no", pool="max") # id_b = ft_netAB(num_classes, stride=1, norm="no", pool="max") dis_a = MsImageDis(3) # discriminator for domain a dis_b = MsImageDis(3) # discriminator for domain a gen_a = gen_a.cuda() gen_b = gen_b.cuda() # id_a = id_a.cuda() # id_b = id_b.cuda() dis_a = dis_a.cuda() dis_b = dis_b.cuda() evaluator = Evaluator(gen_a.enc_content, gen_b.enc_content) metric = DistanceMetric(algorithm=args.dist_metric) # evaluator_s = Evaluator(model_s) # metric_s = DistanceMetric(algorithm=args.dist_metric) start_epoch = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model_t.load_state_dict(checkpoint['model']) # model_discriminator.load_state_dict(checkpoint['model_discriminator']) start_epoch = checkpoint['epoch'] print("=> Start epoch {}".format(start_epoch)) if args.evaluate: metric.train(model_t, train_loader) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) exit() current_margin = args.margin criterion_z_s = CrossEntropyLabelSmooth(num_classes=num_classes, epsilon=0.5).cuda() criterion_att = nn.MSELoss().cuda() criterion_z = CrossEntropyLabelSmooth(num_classes=num_classes, epsilon=0.5).cuda() criterion_I = TripletLoss(margin=current_margin).cuda() criterion_I_s = TripletLoss_s(margin=current_margin).cuda() # criterion_D = nn.CrossEntropyLoss().cuda() # print(args) # # setup id opt # # if args.arch == 'ide': # # ignored_params = list(map(id, model_t.model.fc.parameters() )) + list(map(id, model_t.classifier.parameters() )) # # else: # # ignored_params = list(map(id, id_a.classifier1.parameters())) + list(map(id, id_a.classifier2.parameters())) # # ignored_params_s = list(map(id, id_b.classifier1.parameters())) + list(map(id, id_b.classifier2.parameters())) # # # # base_params = filter(lambda p: id(p) not in ignored_params, id_a.parameters()) # # base_params_s = filter(lambda p: id(p) not in ignored_params_s, id_b.parameters()) # # # # if args.use_adam: # # optimizer_ft = torch.optim.Adam([ # # {'params': filter(lambda p: p.requires_grad,base_params), 'lr': args.lr}, # # {'params': filter(lambda p: p.requires_grad, base_params_s), 'lr': args.lr}, # # {'params': id_a.classifier1.parameters(), 'lr': args.lr}, # # {'params': id_a.classifier2.parameters(), 'lr': args.lr}, # # {'params': id_b.classifier1.parameters(), 'lr': args.lr}, # # {'params': id_b.classifier2.parameters(), 'lr': args.lr}, # # ], # # weight_decay=5e-4) # # else: # # optimizer_ft = torch.optim.SGD([ # # {'params': filter(lambda p: p.requires_grad, base_params), 'lr': args.lr}, # # {'params': filter(lambda p: p.requires_grad, base_params_s), 'lr': args.lr}, # # {'params': model_s.classifier.parameters(), 'lr': args.lr}, # # {'params': model_s.attention_module.parameters(), 'lr': args.lr}, # # {'params': model_t.classifier.parameters(), 'lr': args.lr}, # # {'params': model_t.attention_module.parameters(), 'lr': args.lr}, # # ], # # momentum=0.9, # # weight_decay=5e-4, # # nesterov=True) # # # # id_scheduler = WarmupMultiStepLR(optimizer_ft, args.mile_stone, args.gamma, args.warmup_factor, # # args.warmup_iters, args.warmup_methods) # setup dis and gen dis_opt = torch.optim.Adam([ { 'params': dis_a.parameters(), 'lr': args.dis_lr }, { 'params': dis_b.parameters(), 'lr': args.dis_lr }, ], weight_decay=5e-4) gen_opt = torch.optim.Adam([ { 'params': gen_a.parameters(), 'lr': args.gen_lr }, { 'params': gen_b.parameters(), 'lr': args.gen_lr }, ], weight_decay=5e-4) dis_scheduler = WarmupMultiStepLR(dis_opt, args.mile_stone, args.gamma, args.warmup_factor, args.warmup_iters, args.warmup_methods) gen_scheduler = WarmupMultiStepLR(gen_opt, args.mile_stone, args.gamma, args.warmup_factor, args.warmup_iters, args.warmup_methods) trainer = Trainer(gen_a, gen_b, dis_a, dis_b, criterion_z, criterion_I, criterion_att, trainvallabel, 1, 1, 0.15, 0.05, 5) flag = 1 best_top1 = -1 best_top1_s = -1 # Start training for epoch in range(start_epoch, args.epochs): print("Begin Train") # id_scheduler.step() gen_scheduler.step() dis_scheduler.step() trainer.train(epoch, train_loader, dis_opt, gen_opt) # save_checkpoint( { 'content_a': gen_a.enc_content.state_dict(), 'content_b': gen_b.enc_content.state_dict(), 'style_a': gen_a.enc_style.state_dict(), 'style_b': gen_b.enc_style.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, False, epoch, args.logs_dir, fpath='checkpoint.pth.tar') # if epoch < 200: # continue if not epoch % 1 == 0: continue top1 = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, metric) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'content_a': gen_a.enc_content.state_dict(), 'content_b': gen_b.enc_content.state_dict(), 'style_a': gen_a.enc_style.state_dict(), 'style_b': gen_b.enc_style.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, epoch, args.logs_dir, fpath='checkpoint.pth.tar') # # # # # print('Test with best model_t:') # print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. # format(epoch, top1, best_top1, ' *' if is_best else '')) # # print('Test with best model_s:') # print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. # format(epoch, top1_s, best_top1_s, ' *' if is_best else '')) # # checkpoint = load_checkpoint(osp.join(args.logs_dir,'model_best.pth.tar')) # model_t.load_state_dict(checkpoint['model']) # metric.train(model_t, train_loader) # evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, metric) # # checkpoint_s = load_checkpoint(osp.join(args.logs_dir, 's_model_best.pth.tar')) # model_s.load_state_dict(checkpoint_s['model']) # evaluator_s.evaluate(query_loader_s, gallery_loader_s, dataset.query, dataset.gallery, metric) print(args)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir+'/log')) if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else (256, 128) dataset, num_classes, train_loader, trainvallabel, val_loader, query_loader, gallery_loader, query_loader_rgb, gallery_loader_rgb, query_loader_ir, gallery_loader_ir = \ get_data(args.dataset, args.split, args.data_dir, args.height, args.width, args.batch_size, args.num_instances, args.workers, args.combine_trainval, args.flip_prob, args.padding, args.re_prob) model_t = models.create(args.arch, num_classes=num_classes, num_features=args.features, attention_mode=args.att_mode) model_s = models.create(args.arch, num_classes=num_classes, num_features=args.features, attention_mode=args.att_mode) model_ir = models.create(args.arch, num_classes=num_classes, num_features=args.features, attention_mode=args.att_mode) # print(model) USE_CUDA = torch.cuda.is_available() device = torch.device("cuda:0" if USE_CUDA else "cpu") model_t = nn.DataParallel(model_t, device_ids=[0,1,2]) model_t.to(device) model_s = nn.DataParallel(model_s, device_ids=[0,1,2]) model_s.to(device) model_ir = nn.DataParallel(model_ir, device_ids=[0,1,2]) model_ir.to(device) print(num_classes) #model = model.cuda() #model_discriminator = model_discriminator.cuda() #model_discriminator = nn.DataParallel(model_discriminator, device_ids=[0,1,2]) #model_discriminator.to(device) evaluator = Evaluator(model_t) metric = DistanceMetric(algorithm=args.dist_metric) evaluator_s = Evaluator(model_s) metric_s = DistanceMetric(algorithm=args.dist_metric) evaluator_ir = Evaluator(model_ir) metric_ir = DistanceMetric(algorithm=args.dist_metric) start_epoch = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['model']) model_discriminator.load_state_dict(checkpoint['model_discriminator']) start_epoch = checkpoint['epoch'] print("=> Start epoch {}".format(start_epoch)) if args.evaluate: metric.train(model, train_loader) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) exit() current_margin = args.margin #criterion_z = nn.CrossEntropyLoss().cuda() criterion_z = CrossEntropyLabelSmooth(num_classes= num_classes, epsilon=args.epsilon).cuda() criterion_att = nn.MSELoss().cuda() #criterion_I = TripletLoss(margin= current_margin).cuda() #criterion_I = Circle_Rank_loss(margin_1=args.margin_1, margin_2=args.margin_2, alpha_1=args.alpha_1, alpha_2=args.alpha_2).cuda() criterion_I = Rank_loss(margin_1= args.margin_1, margin_2 =args.margin_2, alpha_1 =args.alpha_1, alpha_2= args.alpha_2).cuda() criterion_t = Triplet(margin=current_margin).cuda() print(args) if args.arch == 'ide': ignored_params = list(map(id, model.model.fc.parameters() )) + list(map(id, model.classifier.parameters() )) else: ignored_params = list(map(id, model_t.module.classifier.parameters())) + list(map(id, model_t.module.attention_module.parameters())) ignored_params_s = list(map(id, model_s.module.classifier.parameters())) + list(map(id, model_s.module.attention_module.parameters())) ignored_params_ir = list(map(id, model_ir.module.classifier.parameters())) + list(map(id, model_ir.module.attention_module.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, model_t.parameters()) base_params_s = filter(lambda p: id(p) not in ignored_params_s, model_s.parameters()) base_params_ir = filter(lambda p: id(p) not in ignored_params_ir, model_ir.parameters()) if args.use_adam: optimizer_ft = torch.optim.Adam([ #print("Ranger") #optimizer_ft = Ranger([ {'params': filter(lambda p: p.requires_grad,base_params), 'lr': args.lr}, {'params': filter(lambda p: p.requires_grad,base_params_s), 'lr':args.lr}, {'params': filter(lambda p: p.requires_grad,base_params_ir), 'lr':args.lr}, {'params': model_t.module.classifier.parameters(), 'lr': args.lr}, {'params': model_t.module.attention_module.parameters(), 'lr': args.lr}, {'params': model_s.module.classifier.parameters(), 'lr': args.lr}, {'params': model_s.module.attention_module.parameters(), 'lr': args.lr}, {'params': model_ir.module.classifier.parameters(), 'lr': args.lr}, {'params': model_ir.module.attention_module.parameters(), 'lr': args.lr}, ], weight_decay=5e-4) #optimizer_discriminator = torch.optim.Adam([ # {'params': model_discriminator.module.model.parameters(), 'lr': args.lr}, # {'params': model_discriminator.module.classifier.parameters(), 'lr': args.lr} # ], # weight_decay=5e-4) else: optimizer_ft = torch.optim.SGD([ {'params': filter(lambda p: p.requires_grad,base_params), 'lr': args.lr}, {'params': model.classifier.parameters(), 'lr': args.lr}, ], momentum=0.9, weight_decay=5e-4, nesterov=True) optimizer_discriminator = torch.optim.SGD([ {'params': model_discriminator.model.parameters(), 'lr': args.lr}, {'params': model_discriminator.classifier.parameters(), 'lr': args.lr}, ], momentum=0.9, weight_decay=5e-4, nesterov=True) print(args.warmup_left) mile_stone = [] mile_stone.append(int(args.warmup_left)) mile_stone.append(int(args.warmup_right)) print(mile_stone) scheduler = WarmupMultiStepLR(optimizer_ft, mile_stone, args.gamma, args.warmup_factor, args.warmup_iters, args.warmup_methods) trainer = Trainer(model_t, model_s, model_ir, criterion_z, criterion_I, criterion_att, criterion_t, trainvallabel, 1, 1 ,args.rgb_w, args.ir_w, 1000) flag = 1 best_top1 = -1 best_top1_s = -1 best_top1_ir = -1 # Start training for epoch in range(start_epoch, args.epochs): scheduler.step() triple_loss, tot_loss = trainer.train(epoch, train_loader, optimizer_ft) save_checkpoint({ 'model': model_t.module.state_dict(), #'model_discriminator': model_discriminator.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, False, epoch, args.logs_dir, fpath='checkpoint.pth.tar') if epoch < 1: continue if not epoch % 10 ==0: continue top1 = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, metric) top1_s = evaluator_s.evaluate(query_loader_rgb, gallery_loader_rgb, dataset.query, dataset.gallery, metric_s) top1_ir = evaluator_ir.evaluate(query_loader_ir, gallery_loader_ir, dataset.query, dataset.gallery, metric_ir) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint({ 'model': model_t.module.state_dict(), #'model_discriminator': model_discriminator.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, epoch, args.logs_dir, fpath='checkpoint.pth.tar') is_best_s = top1_s > best_top1_s best_top1_s = max(top1_s, best_top1_s) save_checkpoint_s({ 'model': model_s.module.state_dict(), 'epoch': epoch+1, 'best_top1': best_top1, }, is_best, epoch, args.logs_dir, fpath='s_checkpoint.pth.tar') is_best_ir = top1_ir > best_top1_ir best_top1_ir = max(top1_ir, best_top1_ir) save_checkpoint_ir({ 'model': model_ir.module.state_dict(), 'epoch': epoch+1, 'best_top1': best_top1, }, is_best, epoch, args.logs_dir, fpath='ir_checkpoint.pth.tar') #print('Test with best model:') #print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. # format(epoch, top1, best_top1, ' *' if is_best else '')) #checkpoint = load_checkpoint(osp.join(args.logs_dir,'model_best.pth.tar')) #model.load_state_dict(checkpoint['model']) #metric.train(model, train_loader) #evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, metric) print(args)
def main(args): # seed if args.seed is not None: np.random.seed(args.seed) torch.manual_seed(args.seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False else: torch.backends.cudnn.benchmark = True if args.logs_dir is None: args.logs_dir = osp.join( f'logs/pcb/{args.dataset}', datetime.datetime.today().strftime('%Y-%m-%d_%H-%M-%S')) else: args.logs_dir = osp.join(f'logs/pcb/{args.dataset}', args.logs_dir) if args.train: os.makedirs(args.logs_dir, exist_ok=True) copy_tree('./reid', args.logs_dir + '/scripts/reid') for script in os.listdir('.'): if script.split('.')[-1] == 'py': dst_file = os.path.join(args.logs_dir, 'scripts', os.path.basename(script)) shutil.copyfile(script, dst_file) sys.stdout = Logger(os.path.join(args.logs_dir, 'log.txt'), ) print('Settings:') print(vars(args)) print('\n') # Create data loaders dataset, num_classes, train_loader, query_loader, gallery_loader, camstyle_loader = \ get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.num_workers, args.combine_trainval, args.crop, args.tracking_icams, args.tracking_fps, args.re, 0, args.camstyle) # Create model model = models.create('pcb', feature_dim=args.feature_dim, num_classes=num_classes, norm=args.norm, dropout=args.dropout, last_stride=args.last_stride) # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: resume_fname = osp.join(f'logs/pcb/{args.dataset}', args.resume, 'model_best.pth.tar') model, start_epoch, best_top1 = checkpoint_loader(model, resume_fname) print("=> Last epoch {} best top1 {:.1%}".format( start_epoch, best_top1)) start_epoch += 1 model = nn.DataParallel(model).cuda() # Criterion criterion = nn.CrossEntropyLoss().cuda() # Optimizer if hasattr(model.module, 'base'): # low learning_rate the base network (aka. ResNet-50) base_param_ids = set(map(id, model.module.base.parameters())) new_params = [ p for p in model.parameters() if id(p) not in base_param_ids ] param_groups = [{ 'params': model.module.base.parameters(), 'lr_mult': 0.1 }, { 'params': new_params, 'lr_mult': 1.0 }] else: param_groups = model.parameters() optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Trainer trainer = Trainer(model, criterion) # Evaluator evaluator = Evaluator(model) if args.train: # Schedule learning rate def adjust_lr(epoch): step_size = args.step_size lr = args.lr * (0.1**(epoch // step_size)) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Draw Curve epoch_s = [] loss_s = [] prec_s = [] eval_epoch_s = [] eval_top1_s = [] # Start training for epoch in range(start_epoch + 1, args.epochs + 1): t0 = time.time() adjust_lr(epoch) # train_loss, train_prec = 0, 0 train_loss, train_prec = trainer.train(epoch, train_loader, optimizer, fix_bn=args.fix_bn) if epoch < args.start_save: continue if epoch % 5 == 0: top1 = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) eval_epoch_s.append(epoch) eval_top1_s.append(top1) else: top1 = 0 is_best = top1 >= best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) epoch_s.append(epoch) loss_s.append(train_loss) prec_s.append(train_prec) draw_curve(os.path.join(args.logs_dir, 'train_curve.jpg'), epoch_s, loss_s, prec_s, eval_epoch_s, None, eval_top1_s) t1 = time.time() t_epoch = t1 - t0 print( '\n * Finished epoch {:3d} top1: {:5.1%} best_eval: {:5.1%} {}\n' .format(epoch, top1, best_top1, ' *' if is_best else '')) print( '*************** Epoch takes time: {:^10.2f} *********************\n' .format(t_epoch)) pass # Final test print('Test with best model:') model, start_epoch, best_top1 = checkpoint_loader( model, osp.join(args.logs_dir, 'model_best.pth.tar')) print("=> Start epoch {} best top1 {:.1%}".format( start_epoch, best_top1)) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) else: print("Test:") evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) pass
def main(args): print(args) np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders if args.big_height is None or args.big_width is None or args.target_height is None or args.target_width is None: args.big_height, args.big_width, args.target_height, args.target_width = ( 256, 256, 224, 224) dataset, num_classes, train_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.big_height, args.big_width, args.target_height, args.target_width, args.batch_size, args.num_instances, args.workers, args.combine_trainval) # Create model model = models.create(args.arch, num_classes=num_classes, num_features=args.features) # Load from checkpoint start_epoch = best = 0 if args.weights: checkpoint = load_checkpoint(args.weights) if args.arch == 'cross_trihard_senet101' or args.arch == 'cross_trihard_se_resnet152': del (checkpoint['last_linear.weight']) del (checkpoint['last_linear.bias']) model.base.load_state_dict(checkpoint) #model.base.load_param(args.weights) elif args.arch == 'cross_trihard_mobilenet': del (checkpoint['state_dict']['module.fc.weight']) del (checkpoint['state_dict']['module.fc.bias']) model_dict = model.state_dict() checkpoint_load = { k.replace('module', 'base'): v for k, v in (checkpoint['state_dict']).items() } model_dict.update(checkpoint_load) model.load_state_dict(model_dict) elif args.arch == 'cross_trihard_shufflenetv2': del (checkpoint['classifier.0.weight']) del (checkpoint['classifier.0.bias']) model.base.load_state_dict(checkpoint) elif args.arch == 'cross_trihard_densenet121': del (checkpoint['classifier.weight']) del (checkpoint['classifier.bias']) pattern = re.compile( r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$' ) for key in list(checkpoint.keys()): res = pattern.match(key) if res: new_key = res.group(1) + res.group(2) checkpoint[new_key] = checkpoint[key] del checkpoint[key] model.base.load_state_dict(checkpoint) elif args.arch == 'cross_trihard_vgg19bn': del (checkpoint['classifier.6.weight']) del (checkpoint['classifier.6.bias']) model.base.load_state_dict(checkpoint) else: del (checkpoint['fc.weight']) del (checkpoint['fc.bias']) model.base.load_state_dict(checkpoint) if args.resume and not args.weights: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] print("=> Start epoch {} best top1 {:.1%}".format(start_epoch, best)) model = nn.DataParallel(model).cuda() # Evaluator evaluator = Evaluator(model) if args.evaluate: print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery) return # Criterion ranking_loss = nn.MarginRankingLoss(margin=args.margin).cuda() criterion = { 'crossentropy': nn.CrossEntropyLoss().cuda(), \ 'trihard': TripletLoss(ranking_loss).cuda() } # Optimizer if hasattr(model.module, 'base'): base_params = [] base_bn_params = [] for name, p in model.module.base.named_parameters(): if 'bn' in name: base_bn_params.append(p) else: base_params.append(p) base_param_ids = set(map(id, model.module.base.parameters())) new_params = [ p for p in model.parameters() if id(p) not in base_param_ids ] param_groups = [{ 'params': base_params, 'lr_mult': 0.1 }, { 'params': base_bn_params, 'lr_mult': 0.1 }, { 'params': new_params, 'lr_mult': args.lr_mult }] else: param_groups = model.parameters() if args.optimizer == 0: optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) else: print('Adam') optimizer = torch.optim.Adam(params=param_groups, lr=args.lr, weight_decay=args.weight_decay) # Trainer trainer = Cross_Trihard_Trainer(model, criterion, metric_loss_weight=args.metric_loss_weight) # Schedule learning rate def adjust_lr(epoch): step_size, step_size2, step_size3 = args.step_size, args.step_size2, args.step_size3 #lr = args.lr * (0.1 ** (epoch // step_size)) if epoch <= step_size: lr = args.lr elif epoch <= step_size2: lr = args.lr * 0.1 elif epoch <= step_size3: lr = args.lr * 0.01 else: lr = args.lr * 0.001 for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) return lr # Start training for epoch in range(start_epoch + 1, args.epochs + 1): lr = adjust_lr(epoch) trainer.train(epoch, train_loader, optimizer, lr, warm_up=True, warm_up_ep=args.warm_up_ep) if epoch % args.epoch_inter == 0 or epoch >= args.dense_evaluate: tmp_res = evaluator.evaluate(test_loader, dataset.query, dataset.gallery) print('tmp_res: ', tmp_res) print('best: ', best) if tmp_res > best and epoch >= args.start_save: best = tmp_res save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch, }, False, fpath=osp.join(args.logs_dir, 'pass%d.pth.tar' % (epoch)))
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Create data loaders assert args.num_instances > 1, "num_instances should be greater than 1" assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) # get source data src_dataset, src_extfeat_loader = \ get_source_data(args.src_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # get target data tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \ get_data(args.tgt_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # Create model # Hacking here to let the classifier be the number of source ids if args.src_dataset == 'dukemtmc': model = models.create(args.arch, num_classes=632, pretrained=False) elif args.src_dataset == 'market1501': model = models.create(args.arch, num_classes=676, pretrained=False) else: raise RuntimeError( 'Please specify the number of classes (ids) of the network.') # Load from checkpoint if args.resume: print( 'Resuming checkpoints from finetuned model on another dataset...\n' ) checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict'], strict=False) else: raise RuntimeWarning('Not using a pre-trained model.') model = nn.DataParallel(model).cuda() # evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) # if args.evaluate: return # Criterion criterion = [ TripletLoss(args.margin, args.num_instances, isAvg=True, use_semi=True).cuda(), TripletLoss(args.margin, args.num_instances, isAvg=True, use_semi=True).cuda(), ] # Optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # training stage transformer on input images normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transformer = T.Compose([ T.Resize((args.height, args.width)), T.RandomHorizontalFlip(), T.ToTensor(), normalizer, T.RandomErasing(probability=0.5, sh=0.2, r1=0.3) ]) # # Start training for iter_n in range(args.iteration): if args.lambda_value == 0: source_features = 0 else: # get source datas' feature source_features, _ = extract_features(model, src_extfeat_loader, print_freq=args.print_freq) # synchronization feature order with src_dataset.train source_features = torch.cat([ source_features[f].unsqueeze(0) for f, _, _ in src_dataset.train ], 0) # extract training images' features print('Iteration {}: Extracting Target Dataset Features...'.format( iter_n + 1)) target_features, tarNames = extract_features( model, tgt_extfeat_loader, print_freq=args.print_freq) # synchronization feature order with dataset.train target_features = torch.cat([ target_features[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval ], 0) target_real_label = np.asarray( [tarNames[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval]) numTarID = len(set(target_real_label)) # calculate distance and rerank result print('Calculating feature distances...') target_features = target_features.numpy() cluster = KMeans(n_clusters=numTarID, n_jobs=8, n_init=1) # select & cluster images as training set of this epochs print('Clustering and labeling...') clusterRes = cluster.fit(target_features) labels, centers = clusterRes.labels_, clusterRes.cluster_centers_ # labels = splitLowconfi(target_features,labels,centers) # num_ids = len(set(labels)) # print('Iteration {} have {} training ids'.format(iter_n+1, num_ids)) # generate new dataset new_dataset = [] for (fname, _, cam), label in zip(tgt_dataset.trainval, labels): # if label==-1: continue # dont need to change codes in trainer.py _parsing_input function and sampler function after add 0 new_dataset.append((fname, label, cam)) print('Iteration {} have {} training images'.format( iter_n + 1, len(new_dataset))) train_loader = DataLoader(Preprocessor(new_dataset, root=tgt_dataset.images_dir, transform=train_transformer), batch_size=args.batch_size, num_workers=4, sampler=RandomIdentitySampler( new_dataset, args.num_instances), pin_memory=True, drop_last=True) # train model with new generated dataset trainer = Trainer(model, criterion) evaluator = Evaluator(model, print_freq=args.print_freq) # Start training for epoch in range(args.epochs): # trainer.train(epoch, remRate=0.2+(0.6/args.iteration)*(1+iter_n)) # to at most 80% trainer.train(epoch, train_loader, optimizer) # test only rank_score = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) #print('co-model:\n') #rank_score = evaluatorB.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) # Evaluate rank_score = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': rank_score.market1501[0], }, True, fpath=osp.join(args.logs_dir, 'adapted.pth.tar')) return (rank_score.map, rank_score.market1501[0])
def main(args): # For fast training. np.random.seed(10000) torch.manual_seed(1000) torch.cuda.manual_seed_all(10000) cudnn.benchmark = True device = torch.device('cuda:' + str(args.gpuid)) torch.cuda.set_device(device) # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) print('logs_dir=', args.logs_dir) # Print logs print(args) # Create data loaders dataset, num_classes, source_train_loader, \ query_loader, gallery_loader = get_data(args.data_dir, args.source, args.target, args.height, args.width, args.batch_size, args.re, args.workers,args) # Create model model = models.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=num_classes) # Invariance learning model # num_tgt = len(dataset.target_train) # model_inv = InvNet(args.features, num_tgt, # beta=args.inv_beta, knn=args.knn, # alpha=args.inv_alpha) # Load from checkpoint start_epoch = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) # model_inv.load_state_dict(checkpoint['state_dict_inv']) start_epoch = checkpoint['epoch'] print("=> Start epoch {} ".format(start_epoch)) # Set model model = nn.DataParallel(model, device_ids=[int(args.gpuid)]) # model_inv = model_inv.to(device) # Evaluator evaluator = Evaluator(model) if args.evaluate: print("Test:") evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, args.print_freq, args.output_feature) return # Optimizer base_param_ids = set(map(id, model.module.base.parameters())) base_params_need_for_grad = filter(lambda p: p.requires_grad, model.module.base.parameters()) new_params = [p for p in model.parameters() if id(p) not in base_param_ids] param_groups = [{ 'params': base_params_need_for_grad, 'lr_mult': 0.1 }, { 'params': new_params, 'lr_mult': 1.0 }] optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Trainer # trainer = Trainer(model, '', lmd=args.lmd) trainer = Trainer(model, None, lmd=args.lmd) # Schedule learning rate def adjust_lr(epoch): step_size = args.epochs_decay lr = args.lr * (0.1**(epoch // step_size)) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): print(epoch) adjust_lr(epoch) trainer.sys_train(epoch, source_train_loader, optimizer, print_freq=args.print_freq) save_checkpoint( { 'state_dict': model.module.state_dict(), # 'state_dict_inv': model_inv.state_dict(), 'epoch': epoch + 1, }, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} \n'.format(epoch)) if epoch % 5 == 0: print(' test with epoch {:3d} : ---'.format(epoch)) evaluator = Evaluator(model) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, args.print_freq, args.output_feature) # Final test print('Final Test : Testing with best model......') evaluator = Evaluator(model) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, args.print_freq, args.output_feature)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir + '/log')) if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else (256, 128) dataset, num_classes, train_loader, trainvallabel, val_loader, query_loader, gallery_loader, query_loader_s, gallery_loader_s = \ get_data(args.dataset, args.split, args.data_dir, args.height, args.width, args.batch_size, args.num_instances, args.workers, args.combine_trainval, args.flip_prob, args.padding, args.re_prob) print(num_classes) # model_s, model_t, model_discriminator = models.create(args.arch, num_classes=num_classes, num_features=args.features) model_t = models.create(args.arch, num_classes=num_classes, num_features=args.features, attention_mode=args.att_mode) model_s = models.create(args.arch, num_classes=num_classes, num_features=args.features, attention_mode=args.att_mode) # load source network # checkpoint_s = load_checkpoint('/home/fan/cross_reid/source_net/model_best.pth.tar') # # model_dict = model_s.state_dict() # state_dict = {k:v for k,v in checkpoint_s.items() if k in model_dict.keys()} # model_dict.update(state_dict) # print(model_s) # print(checkpoint_s['model']) # print(model_dict.keys()) # model_s.load_state_dict(model_dict) # print(model_s) USE_CUDA = torch.cuda.is_available() device = torch.device("cuda:0" if USE_CUDA else "cpu") model_s = nn.DataParallel(model_s, device_ids=[0, 2]) model_s.to(device) # model_s = model_s.cuda() model_t = nn.DataParallel(model_t, device_ids=[0, 2]) model_t.to(device) # model_discriminator = model_discriminator.cuda() evaluator = Evaluator(model_t) metric = DistanceMetric(algorithm=args.dist_metric) evaluator_s = Evaluator(model_s) metric_s = DistanceMetric(algorithm=args.dist_metric) start_epoch = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model_t.load_state_dict(checkpoint['model']) # model_discriminator.load_state_dict(checkpoint['model_discriminator']) start_epoch = checkpoint['epoch'] print("=> Start epoch {}".format(start_epoch)) if args.evaluate: metric.train(model_t, train_loader) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) exit() current_margin = args.margin criterion_z_s = CrossEntropyLabelSmooth(num_classes=num_classes, epsilon=0.5).cuda() criterion_att = nn.MSELoss().cuda() criterion_z = CrossEntropyLabelSmooth(num_classes=num_classes, epsilon=0.5).cuda() criterion_I = TripletLoss(margin=current_margin).cuda() criterion_I_s = TripletLoss_s(margin=current_margin).cuda() # criterion_D = nn.CrossEntropyLoss().cuda() print(args) if args.arch == 'ide': ignored_params = list(map(id, model_t.model.fc.parameters())) + list( map(id, model_t.classifier.parameters())) else: ignored_params = list(map( id, model_t.module.classifier.parameters())) + list( map(id, model_t.module.attention_module.parameters())) ignored_params_s = list(map( id, model_s.module.classifier.parameters())) + list( map(id, model_s.module.attention_module.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, model_t.parameters()) base_params_s = filter(lambda p: id(p) not in ignored_params_s, model_s.parameters()) if args.use_adam: optimizer_ft = torch.optim.Adam([ { 'params': filter(lambda p: p.requires_grad, base_params), 'lr': args.lr }, { 'params': filter(lambda p: p.requires_grad, base_params_s), 'lr': args.lr }, { 'params': model_s.module.classifier.parameters(), 'lr': args.lr }, { 'params': model_s.module.attention_module.parameters(), 'lr': args.lr }, { 'params': model_t.module.classifier.parameters(), 'lr': args.lr }, { 'params': model_t.module.attention_module.parameters(), 'lr': args.lr }, ], weight_decay=5e-4) # optimizer_discriminator = torch.optim.Adam([ # {'params': model_discriminator.model.parameters(), 'lr': args.lr}, # {'params': model_discriminator.classifier.parameters(), 'lr': args.lr} # ], # weight_decay=5e-4) else: optimizer_ft = torch.optim.SGD([ { 'params': filter(lambda p: p.requires_grad, base_params), 'lr': args.lr }, { 'params': filter(lambda p: p.requires_grad, base_params_s), 'lr': args.lr }, { 'params': model_s.classifier.parameters(), 'lr': args.lr }, { 'params': model_s.attention_module.parameters(), 'lr': args.lr }, { 'params': model_t.classifier.parameters(), 'lr': args.lr }, { 'params': model_t.attention_module.parameters(), 'lr': args.lr }, ], momentum=0.9, weight_decay=5e-4, nesterov=True) # optimizer_discriminator = torch.optim.SGD([ # {'params': model_discriminator.model.parameters(), 'lr': args.lr}, # {'params': model_discriminator.classifier.parameters(), 'lr': args.lr}, # ], # momentum=0.9, # weight_decay=5e-4, # nesterov=True) scheduler = WarmupMultiStepLR(optimizer_ft, args.mile_stone, args.gamma, args.warmup_factor, args.warmup_iters, args.warmup_methods) trainer = Trainer(model_s, model_t, criterion_z, criterion_z_s, criterion_I, criterion_I_s, criterion_att, trainvallabel, 1, 1, 0.15, 0.05, 5) flag = 1 best_top1 = -1 best_top1_s = -1 # Start training for epoch in range(start_epoch, args.epochs): scheduler.step() triple_loss, tot_loss, att_loss = trainer.train( epoch, train_loader, optimizer_ft) save_checkpoint( { 'model': model_t.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, False, epoch, args.logs_dir, fpath='checkpoint.pth.tar') # checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) # for i, keys in enumerate(checkpoint): # # model_t.load_state_dict(checkpoint['model']) # print(keys) # print(i) # if epoch < 200: # continue if not epoch % 10 == 0: continue top1 = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, metric) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'model': model_t.module.state_dict(), # 'model_discriminator': model_discriminator.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, epoch, args.logs_dir, fpath='checkpoint.pth.tar') top1_s = evaluator_s.evaluate(query_loader_s, gallery_loader_s, dataset.query, dataset.gallery, metric) is_best_s = top1_s > best_top1_s best_top1_s = max(top1_s, best_top1_s) save_checkpoint_s( { 'model': model_s.module.state_dict(), # 'model_discriminator': model_discriminator.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1_s, }, is_best, epoch, args.logs_dir, fpath='s_checkpoint.pth.tar') print('Test with best model_t:') print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'.format( epoch, top1, best_top1, ' *' if is_best else '')) print('Test with best model_s:') print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'.format( epoch, top1_s, best_top1_s, ' *' if is_best else '')) # state_dict = {k:v for k,v in checkpoint_s.items() if k in model_dict.keys()} # model_dict.update(state_dict) metric.train(model_t, train_loader) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, metric) checkpoint_s = load_checkpoint( osp.join(args.logs_dir, 's_model_best.pth.tar')) model_s.load_state_dict(checkpoint_s['model']) evaluator_s.evaluate(query_loader_s, gallery_loader_s, dataset.query, dataset.gallery, metric) print(args)
def main(args): os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders assert args.num_instances > 1, "num_instances should be greater than 1" assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) ## get_source_data src_dataset, src_extfeat_loader = \ get_source_data(args.src_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # get_target_data tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \ get_data(args.tgt_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # Create model # Hacking here to let the classifier be the last feature embedding layer # Net structure: avgpool -> FC(2048) -> FC(args.features) if args.src_dataset == 'dukemtmc': model = models.create(args.arch, num_classes=0, num_split=args.num_split, cluster=args.dce_loss) #duke elif args.src_dataset == 'market1501': model = models.create(args.arch, num_classes=0, num_split=args.num_split, cluster=args.dce_loss) else: raise RuntimeError( 'Please specify the number of classes (ids) of the network.') # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: print( 'Resuming checkpoints from finetuned model on another dataset...\n' ) checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint, strict=False) else: raise RuntimeWarning('Not using a pre-trained model') model = nn.DataParallel(model).cuda() # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model, print_freq=args.print_freq) print("Test with the original model trained on source domain:") best_top1 = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) if args.evaluate: return # Criterion criterion = [] criterion.append( TripletLoss(margin=args.margin, num_instances=args.num_instances).cuda()) criterion.append( TripletLoss(margin=args.margin, num_instances=args.num_instances).cuda()) #multi lr base_param_ids = set(map(id, model.module.base.parameters())) new_params = [p for p in model.parameters() if id(p) not in base_param_ids] param_groups = [{ 'params': model.module.base.parameters(), 'lr_mult': 1.0 }, { 'params': new_params, 'lr_mult': 1.0 }] # Optimizer optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=0.9, weight_decay=args.weight_decay) ##### adjust lr def adjust_lr(epoch): if epoch <= 7: lr = args.lr elif epoch <= 14: lr = 0.3 * args.lr else: lr = 0.1 * args.lr for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) ##### training stage transformer on input images normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transformer = T.Compose([ Resize((args.height, args.width)), T.RandomHorizontalFlip(), T.ToTensor(), normalizer, T.RandomErasing(probability=0.5, sh=0.2, r1=0.3) ]) # Start training iter_nums = args.iteration cluster_list = [] top_percent = args.rho for iter_n in range(0, iter_nums): #### get source datas' feature if args.load_dist and iter_n == 0: dist = pickle.load( open('dist' + str(args.num_split) + '.pkl', 'rb')) euclidean_dist_list = dist['euclidean'] rerank_dist_list = dist['rerank'] else: source_features, _ = extract_features(model, src_extfeat_loader, for_eval=False) if isinstance(source_features[src_dataset.trainval[0][0]], list): len_f = len(source_features[src_dataset.trainval[0][0]]) source_features = [ torch.cat([ source_features[f][i].unsqueeze(0) for f, _, _ in src_dataset.trainval ], 0) for i in range(len_f) ] else: source_features = torch.cat([ source_features[f].unsqueeze(0) for f, _, _ in src_dataset.trainval ], 0) # synchronization feature order with s_dataset.trainval #### extract training images' features print('Iteration {}: Extracting Target Dataset Features...'.format( iter_n + 1)) target_features, _ = extract_features(model, tgt_extfeat_loader, for_eval=False) if isinstance(target_features[tgt_dataset.trainval[0][0]], list): len_f = len(target_features[tgt_dataset.trainval[0][0]]) target_features = [ torch.cat([ target_features[f][i].unsqueeze(0) for f, _, _ in tgt_dataset.trainval ], 0) for i in range(len_f) ] else: target_features = torch.cat([ target_features[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval ], 0) # synchronization feature order with dataset.trainval #### calculate distance and rerank result print('Calculating feature distances...') # target_features = target_features.numpy() euclidean_dist_list, rerank_dist_list = compute_dist( source_features, target_features, lambda_value=args.lambda_value, no_rerank=args.no_rerank, num_split=args.num_split) # lambda=1 means only source dist del target_features del source_features labels_list, cluster_list = generate_selflabel(euclidean_dist_list, rerank_dist_list, iter_n, args, cluster_list) #### generate new dataset train_loader_list = generate_dataloader(tgt_dataset, labels_list, train_transformer, iter_n, args) del labels_list # del cluster_list top1 = iter_trainer(model, tgt_dataset, train_loader_list, test_loader, optimizer, criterion, args.epochs, args.logs_dir, args.print_freq) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': iter_n + 1, 'best_top1': best_top1, # 'num_ids': num_ids, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(iter_n + 1, top1, best_top1, ' *' if is_best else ''))
def run(): np.random.seed(opt.seed) torch.manual_seed(opt.seed) cudnn.benchmark = True data_dir = opt.data_dir # Redirect print to both console and log file #if not opt.evaluate: # sys.stdout = Logger(osp.join(opt.logs_dir, 'log_l2_per.txt')) # Create data loaders def readlist(path): lines = [] with open(path, 'r') as f: data = f.readlines() #pdb.set_trace() for line in data: name, pid, cam = line.split() lines.append((name, int(pid), int(cam))) return lines # Load data list for wuzhen if osp.exists(osp.join(data_dir, 'train.txt')): train_list = readlist(osp.join(data_dir, 'train.txt')) else: print("The training list doesn't exist") if osp.exists(osp.join(data_dir, 'val.txt')): val_list = readlist(osp.join(data_dir, 'val.txt')) else: print("The validation list doesn't exist") if osp.exists(osp.join(data_dir, 'query.txt')): query_list = readlist(osp.join(data_dir, 'query.txt')) else: print("The query.txt doesn't exist") if osp.exists(osp.join(data_dir, 'gallery.txt')): gallery_list = readlist(osp.join(data_dir, 'gallery.txt')) else: print("The gallery.txt doesn't exist") if opt.height is None or opt.width is None: opt.height, opt.width = (144, 56) if opt.arch == 'inception' else \ (256, 128) train_loader,val_loader, test_loader = \ get_data(opt.split, data_dir, opt.height, opt.width, opt.batchSize, opt.workers, opt.combine_trainval, train_list, val_list, query_list, gallery_list) # Create model # ori 14514; clear 12654, 16645 densenet = densenet121(num_classes=20330, num_features=256) start_epoch = best_top1 = 0 if opt.resume: #checkpoint = load_checkpoint(opt.resume) #densenet.load_state_dict(checkpoint['state_dict']) densenet.load_state_dict(torch.load(opt.resume)) start_epoch = opt.resume_epoch print("=> Finetune Start epoch {} ".format(start_epoch)) if opt.pretrained_model: print('Start load params...') load_params(densenet, opt.pretrained_model) # Load from checkpoint #densenet = nn.DataParallel(densenet).cuda() metric = DistanceMetric(algorithm=opt.dist_metric) print('densenet') show_info(densenet, with_arch=True, with_grad=False) netG = netg() print('netG') show_info(netG, with_arch=True, with_grad=False) netG.apply(weights_init) if opt.netG != '': netG.load_state_dict(torch.load(opt.netG)) #load_params(netG,opt.netG) if opt.cuda: netG = netG.cuda() densenet = densenet.cuda() perceptionloss = perception_loss(cuda=opt.cuda) l2loss = l2_loss(cuda=opt.cuda) # discriloss=discri_loss(cuda = opt.cuda,batchsize = opt.batchSize,height = \ # opt.height,width = opt.width,lr = opt.lr,step_size = \ # opt.step_size,decay_step = opt.decay_step ) # Evaluator evaluator = Evaluator(densenet) # if opt.evaluate: metric.train(densenet, train_loader) print("Validation:") evaluator.evaluate(val_loader, val_list, val_list, metric) print("Test:") evaluator.evaluate(test_loader, query_list, gallery_list, metric) # return # Criterion # criterion = nn.CrossEntropyLoss(ignore_index=-100).cuda() criterion = nn.CrossEntropyLoss().cuda() # Optimizer param_groups = [] mult_lr(densenet, param_groups) optimizer = optim.SGD(param_groups, lr=opt.lr, momentum=opt.momentum, weight_decay=opt.weight_decay) # optimizer = optim.Adam(param_groups, lr=opt.lr, betas=(opt.beta1, 0.9)) optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.9)) # Start training for epoch in range(start_epoch, opt.epochs): adjust_lr(optimizer, epoch) adjust_lr(optimizerG, epoch) #discriloss.adjust_lr(epoch) losses = AverageMeter() precisions = AverageMeter() densenet.train() for i, data in enumerate(train_loader): ############################ # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) ########################### # train with real real_cpu, _, pids, _ = data if opt.cuda: real_cpu = real_cpu.cuda() targets = Variable(pids.cuda()) input.resize_as_(real_cpu).copy_(real_cpu) inputv = Variable(input) outputs, output_dense, _ = densenet(inputv) fake = netG(output_dense) fake = fake * 3 #discriloss(fake = fake, inputv = inputv, i = i) ############################ # (2) Update G network: maximize log(D(G(z))) ########################### if i % opt.CRITIC_ITERS == 0: netG.zero_grad() optimizer.zero_grad() #loss_discri = discriloss.gloss(fake = fake) loss_l2 = l2loss(fake=fake, inputv=inputv) loss_perception = perceptionloss(fake=fake, inputv=inputv) loss_classify = criterion(outputs, targets) prec, = accuracy(outputs.data, targets.data) prec = prec[0] losses.update(loss_classify.data[0], targets.size(0)) precisions.update(prec, targets.size(0)) loss = loss_classify + 0 * loss_l2 + 0 * loss_perception # loss = loss_discri loss.backward() optimizerG.step() optimizer.step() #print(precisions.val) #print(precisions.avg) # print('[%d/%d][%d/%d] '%(epoch, opt.epochs, i, len(train_loader))) # print('[%d/%d][%d/%d] Loss_discri: %.4f '%(epoch, opt.epochs, i, \ # len(train_loader),loss_discri.data[0])) print('[%d/%d][%d/%d] Loss_l2: %.4f Loss_perception: %.4f '%(epoch, opt.epochs, i, \ len(train_loader),loss_l2.data[0],loss_perception.data[0])) print('Loss {}({})\t' 'Prec {}({})\t'.format(losses.val, losses.avg, precisions.val, precisions.avg)) if i % 100 == 0: vutils.save_image(real_cpu, '%s/real_samples.png' % opt.outf, normalize=True) outputs, output_dense, _ = densenet(x=inputv) fake = netG(output_dense) fake = fake * 3 vutils.save_image(fake.data, '%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch), normalize=True) show_info(densenet, with_arch=False, with_grad=True) show_info(netG, with_arch=False, with_grad=True) if epoch % 5 == 0: torch.save(densenet.state_dict(), '%s/densenet_epoch_%d.pth' % (opt.outf, epoch)) torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (opt.outf, epoch)) if epoch < opt.start_save: continue top1 = evaluator.evaluate(val_loader, val_list, val_list) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'state_dict': densenet.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(opt.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(epoch, top1, best_top1, ' *' if is_best else '')) if (epoch + 1) % 5 == 0: print('Test model: \n') evaluator.evaluate(test_loader, query_list, gallery_list) model_name = 'epoch_' + str(epoch) + '.pth.tar' torch.save({'state_dict': densenet.state_dict()}, osp.join(opt.logs_dir, model_name)) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(opt.logs_dir, 'model_best.pth.tar')) densenet.load_state_dict(checkpoint['state_dict']) print('best epoch: ', checkpoint['epoch']) metric.train(densenet, train_loader) evaluator.evaluate(test_loader, query_list, gallery_list, metric)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Create data loaders assert args.num_instances > 1, "num_instances should be greater than 1" assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) # get source data src_dataset, src_extfeat_loader = \ get_source_data(args.src_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # get target data tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \ get_data(args.tgt_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # Create model # Hacking here to let the classifier be the number of source ids if args.src_dataset == 'dukemtmc': model = models.create(args.arch, num_classes=632, pretrained=False) elif args.src_dataset == 'market1501': model = models.create(args.arch, num_classes=676, pretrained=False) else: raise RuntimeError( 'Please specify the number of classes (ids) of the network.') # Load from checkpoint if args.resume: print( 'Resuming checkpoints from finetuned model on another dataset...\n' ) checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict'], strict=False) else: raise RuntimeWarning('Not using a pre-trained model.') model = nn.DataParallel(model).cuda() # evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) # if args.evaluate: return # Criterion criterion = [ TripletLoss(args.margin, args.num_instances, isAvg=True, use_semi=True).cuda(), TripletLoss(args.margin, args.num_instances, isAvg=True, use_semi=True).cuda(), None, None ] # Optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # training stage transformer on input images normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transformer = T.Compose([ T.Resize((args.height, args.width)), T.RandomHorizontalFlip(), T.ToTensor(), normalizer, T.RandomErasing(probability=0.5, sh=0.2, r1=0.3) ]) evaluator = Evaluator(model, print_freq=args.print_freq) evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) # # Start training for iter_n in range(args.iteration): if args.lambda_value == 0: source_features = 0 else: # get source datas' feature source_features, _ = extract_features(model, src_extfeat_loader, print_freq=args.print_freq) # synchronization feature order with src_dataset.train source_features = torch.cat([ source_features[f].unsqueeze(0) for f, _, _, _ in src_dataset.train ], 0) # extract training images' features print('Iteration {}: Extracting Target Dataset Features...'.format( iter_n + 1)) target_features, tarNames = extract_features( model, tgt_extfeat_loader, print_freq=args.print_freq) # synchronization feature order with dataset.train target_features = torch.cat([ target_features[f].unsqueeze(0) for f, _, _, _ in tgt_dataset.trainval ], 0) # target_real_label = np.asarray([tarNames[f].unsqueeze(0) for f, _, _, _ in tgt_dataset.trainval]) # calculate distance and rerank result # method 1 target_features = target_features.numpy() rerank_dist = re_ranking(source_features, target_features, lambda_value=args.lambda_value) # method 2 # distmat_qq = calDis(source_features, source_features) # distmat_qg = calDis(source_features, target_features) # distmat_gg = calDis(target_features, target_features) # rerank_dist = re_ranking2(distmat_qg.numpy(), distmat_qq.numpy(), distmat_gg.numpy()) cluster = HDBSCAN(metric='precomputed', min_samples=10) # select & cluster images as training set of this epochs clusterRes = cluster.fit(rerank_dist) labels, label_num = clusterRes.labels_, clusterRes.labels_.max() + 1 centers = np.zeros((label_num, target_features.shape[1])) nums = [0] * target_features.shape[1] print('clusters num =', label_num) # generate new dataset new_dataset = [] index = -1 for (fname, _, cam, timestamp), label in zip(tgt_dataset.trainval, labels): index += 1 if label == -1: continue # dont need to change codes in trainer.py _parsing_input function and sampler function after add 0 new_dataset.append((fname, label, cam, timestamp)) centers[label] += target_features[index] nums[label] += 1 print('Iteration {} have {} training images'.format( iter_n + 1, len(new_dataset))) train_loader = DataLoader(Preprocessor(new_dataset, root=tgt_dataset.images_dir, transform=train_transformer), batch_size=args.batch_size, num_workers=4, sampler=RandomIdentitySampler( new_dataset, args.num_instances), pin_memory=True, drop_last=True) for i in range(label_num): centers[i] /= nums[i] criterion[3] = ClassificationLoss(normalize(centers, axis=1)).cuda() classOptimizer = torch.optim.Adam( [{ 'params': model.parameters() }, { 'params': criterion[3].classifier.parameters(), 'lr': 1e-3 }], lr=args.lr) class_trainer = ClassificationTrainer(model, train_loader, criterion, classOptimizer) for epoch in range(args.epochs): class_trainer.train(epoch) rank_score = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) # Evaluate rank_score = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': rank_score.market1501[0], }, True, fpath=osp.join(args.logs_dir, 'adapted.pth.tar')) return (rank_score.map, rank_score.market1501[0])
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders assert args.num_instances > 1, "num_instances should be greater than 1" assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) dataset, num_classes, train_loader, val_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.height, args.width, args.batch_size, args.num_instances, args.workers, args.combine_trainval) # Create model # Hacking here to let the classifier be the last feature embedding layer # Net structure: avgpool -> FC(1024) -> FC(args.features) model = models.create(args.arch, num_features=1024, dropout=args.dropout, num_classes=args.features) # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] best_top1 = checkpoint['best_top1'] print("=> Start epoch {} best top1 {:.1%}" .format(start_epoch, best_top1)) model = nn.DataParallel(model).cuda() # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model) if args.evaluate: metric.train(model, train_loader) print("Validation:") evaluator.evaluate(val_loader, dataset.val, dataset.val, metric) print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric) return # Criterion criterion = TripletLoss(margin=args.margin).cuda() # Optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # Trainer trainer = Trainer(model, criterion) # Schedule learning rate def adjust_lr(epoch): lr = args.lr if epoch <= 100 else \ args.lr * (0.001 ** ((epoch - 100) / 50.0)) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, train_loader, optimizer) if epoch < args.start_save: continue top1 = evaluator.evaluate(val_loader, dataset.val, dataset.val) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint({ 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(epoch, top1, best_top1, ' *' if is_best else '')) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) model.module.load_state_dict(checkpoint['state_dict']) metric.train(model, train_loader) evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
def main(args): # seed if args.seed is not None: np.random.seed(args.seed) torch.manual_seed(args.seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False else: torch.backends.cudnn.benchmark = True if args.logs_dir is None: args.logs_dir = osp.join( f'logs/triplet/{args.dataset}', datetime.datetime.today().strftime('%Y-%m-%d_%H-%M-%S')) else: args.logs_dir = osp.join(f'logs/triplet/{args.dataset}', args.logs_dir) if args.train: os.makedirs(args.logs_dir, exist_ok=True) copy_tree('./reid', args.logs_dir + '/scripts/reid') for script in os.listdir('.'): if script.split('.')[-1] == 'py': dst_file = os.path.join(args.logs_dir, 'scripts', os.path.basename(script)) shutil.copyfile(script, dst_file) sys.stdout = Logger(os.path.join(args.logs_dir, 'log.txt'), ) print('Settings:') print(vars(args)) print('\n') # Create data loaders assert args.num_instances > 1, "num_instances should be larger than 1" assert args.batch_size % args.num_instances == 0, 'num_instances should divide batch_size' dataset, num_classes, train_loader, query_loader, gallery_loader, _ = \ get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.num_workers, args.combine_trainval, args.crop, args.tracking_icams, args.tracking_fps, args.re, args.num_instances, False) # Create model for triplet (num_classes = 0, num_instances > 0) model = models.create('ide', feature_dim=args.feature_dim, num_classes=0, norm=args.norm, dropout=args.dropout, last_stride=args.last_stride) # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: resume_fname = osp.join(f'logs/triplet/{args.dataset}', args.resume, 'model_best.pth.tar') model, start_epoch, best_top1 = checkpoint_loader(model, resume_fname) print("=> Last epoch {} best top1 {:.1%}".format( start_epoch, best_top1)) start_epoch += 1 model = nn.DataParallel(model).cuda() # Criterion criterion = TripletLoss(margin=args.margin).cuda() # Optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # Trainer trainer = Trainer(model, criterion) # Evaluator evaluator = Evaluator(model) if args.train: # Schedule learning rate def adjust_lr(epoch): if epoch <= args.step_size: lr = args.lr else: lr = args.lr * (0.001**(float(epoch - args.step_size) / (args.epochs - args.step_size))) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Draw Curve epoch_s = [] loss_s = [] prec_s = [] eval_epoch_s = [] eval_top1_s = [] # Start training for epoch in range(start_epoch + 1, args.epochs + 1): adjust_lr(epoch) # train_loss, train_prec = 0, 0 train_loss, train_prec = trainer.train(epoch, train_loader, optimizer, fix_bn=args.fix_bn) if epoch < args.start_save: continue if epoch % 25 == 0: top1 = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) eval_epoch_s.append(epoch) eval_top1_s.append(top1) else: top1 = 0 is_best = top1 >= best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) epoch_s.append(epoch) loss_s.append(train_loss) prec_s.append(train_prec) draw_curve(os.path.join(args.logs_dir, 'train_curve.jpg'), epoch_s, loss_s, prec_s, eval_epoch_s, None, eval_top1_s) pass # Final test print('Test with best model:') model, start_epoch, best_top1 = checkpoint_loader( model, osp.join(args.logs_dir, 'model_best.pth.tar')) print("=> Start epoch {} best top1 {:.1%}".format( start_epoch, best_top1)) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) else: print("Test:") evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) pass
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders if args.loss == 'triplet': assert args.num_instances > 1, 'TripletLoss requires num_instances > 1' assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' dataset, num_classes, train_loader, val_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.batch_size, args.workers, args.num_instances, combine_trainval=args.combine_trainval) # Create model if args.loss == 'xentropy': model = InceptionNet(num_classes=num_classes, num_features=args.features, dropout=args.dropout) elif args.loss == 'oim': model = InceptionNet(num_features=args.features, norm=True, dropout=args.dropout) elif args.loss == 'triplet': model = InceptionNet(num_features=args.features, dropout=args.dropout) else: raise ValueError("Cannot recognize loss type:", args.loss) model = torch.nn.DataParallel(model).cuda() # Load from checkpoint if args.resume: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] best_top1 = checkpoint['best_top1'] print("=> start epoch {} best top1 {:.1%}".format( args.start_epoch, best_top1)) else: best_top1 = 0 # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model) if args.evaluate: metric.train(model, train_loader) print("Validation:") evaluator.evaluate(val_loader, dataset.val, dataset.val, metric) print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric) return # Criterion if args.loss == 'xentropy': criterion = torch.nn.CrossEntropyLoss() elif args.loss == 'oim': criterion = OIMLoss(model.module.num_features, num_classes, scalar=args.oim_scalar, momentum=args.oim_momentum) elif args.loss == 'triplet': criterion = TripletLoss(margin=args.triplet_margin) else: raise ValueError("Cannot recognize loss type:", args.loss) criterion.cuda() # Optimizer if args.optimizer == 'sgd': optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) elif args.optimizer == 'adam': optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) else: raise ValueError("Cannot recognize optimizer type:", args.optimizer) # Trainer trainer = Trainer(model, criterion) # Schedule learning rate def adjust_lr(epoch): if args.optimizer == 'sgd': lr = args.lr * (0.1**(epoch // 60)) elif args.optimizer == 'adam': lr = args.lr if epoch <= 100 else \ args.lr * (0.001 ** (epoch - 100) / 50) else: raise ValueError("Cannot recognize optimizer type:", args.optimizer) for g in optimizer.param_groups: g['lr'] = lr # Start training for epoch in range(args.start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, train_loader, optimizer) top1 = evaluator.evaluate(val_loader, dataset.val, dataset.val) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'state_dict': model.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(epoch, top1, best_top1, ' *' if is_best else '')) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) model.load_state_dict(checkpoint['state_dict']) metric.train(model, train_loader) evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True logger = TensorLogger(osp.join(args.log_root, 'Tensorboard_logs', args.logs_dir)) # Redirect print to both console and log file logs_dir = osp.join(args.log_root, args.logs_dir) if not args.evaluate: sys.stdout = Logger(osp.join(logs_dir, 'log.txt')) # Create data loaders # assert args.num_instances > 1, "num_instances should be greater than 1" assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) dataset, num_classes, train_loader, val_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.height, args.width, args.batch_size, args.num_instances, args.workers, args.combine_trainval) print("lr:", args.lr) print("max_epoch:", args.epochs) # Create model # Hacking here to let the classifier be the last feature embedding layer # Net structure: avgpool -> FC(1024) -> FC(args.features) model = models.create(args.arch, num_features=1024, num_diff_features=128, dropout=args.dropout, cut_at_pooling=False) print(model) # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: checkpoint = load_checkpoint(args.resume) pretrained_dict = {k: v for k, v in checkpoint['state_dict'].items() if k in model.state_dict()} model_dict = model.state_dict() model_dict.update(pretrained_dict) model.load_state_dict(model_dict) start_epoch = checkpoint['epoch'] best_top1 = checkpoint['best_top1'] print("=> Start epoch {} best top1 {:.1%}" .format(start_epoch, best_top1)) model = nn.DataParallel(model).cuda() # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model) if args.evaluate: metric.train(model, train_loader) # print("Validation:") # evaluator.evaluate(val_loader, dataset.val, dataset.val, metric) print("Test:") if len(args.dataset)>1: for dataset_name in args.dataset: print("{} test result:".format(dataset_name)) evaluator.evaluate(test_loader[dataset_name], dataset.query[dataset_name], dataset.gallery[dataset_name], metric) return else: evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric) return # Criterion criterion = TripletLoss(margin=args.margin).cuda() # Optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # Trainer trainer = Trainer(model, criterion) # Schedule learning rate def adjust_lr(epoch): # lr = args.lr if epoch <= 270 else \ # args.lr * (0.001 ** ((epoch - 270) / 135)) # lr = args.lr if epoch <= 100 else \ # args.lr * (0.001 ** ((epoch - 100) / 50.0)) lr = args.lr if epoch <= args.lr_change_epochs else \ args.lr * (0.001 ** ((epoch - args.lr_change_epochs)/float(args.epochs-args.lr_change_epochs))) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # epoch_num = args.maxstep//(750//18) + 1 # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, train_loader, optimizer, logger) #######Tensorboard-logs########## for tag, value in model.named_parameters(): tag = tag.replace('.', '/') try: logger.histo_summary(tag, to_np(value), epoch*len(train_loader) + 1) logger.histo_summary(tag + '/grad', to_np(value.grad), epoch*len(train_loader) + 1) except AttributeError: pass # top1 = evaluator.evaluate(val_loader, dataset.val, dataset.val) # top1 = evaluator.evaluate(test_loader, dataset.query, dataset.gallery) top1 = 1 # is_best = top1 > best_top1 # best_top1 = max(top1, best_top1) # save_checkpoint({ # 'state_dict': model.module.state_dict(), # 'epoch': epoch + 1, # 'best_top1': best_top1, # }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) is_best = False best_top1 = 1 save_checkpoint({ 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(epoch, top1, best_top1, ' *' if is_best else '')) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(logs_dir, 'checkpoint.pth.tar')) model.module.load_state_dict(checkpoint['state_dict']) metric.train(model, train_loader) print("Test:") if len(args.dataset) > 1: for dataset_name in args.dataset: print("{} test result:".format(dataset_name)) evaluator.evaluate(test_loader[dataset_name], dataset.query[dataset_name], dataset.gallery[dataset_name], metric) else: evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) print(args) # Create data loaders if args.big_height is None or args.big_width is None or args.target_height is None or args.target_width is None: args.big_height, args.big_width, args.target_height, args.target_width = ( 256, 256, 224, 224) dataset, num_classes, train_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.big_height, args.big_width, args.target_height, args.target_width, args.batch_size, args.num_instances, args.workers, args.combine_trainval) # Create models model = models.create(name=args.arch, num_classes=num_classes, num_features=args.features, norm=True) print(model) # Load from checkpoint start_epoch = best = 0 if args.weights and hasattr(model, 'base'): print('loading resnet50') checkpoint = load_checkpoint(args.weights) del (checkpoint['fc.weight']) del (checkpoint['fc.bias']) model.base.load_state_dict(checkpoint) if args.resume and not args.weights: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] print("=> Start epoch {} best top1 {:.1%}".format(start_epoch, best)) model = nn.DataParallel(model).cuda() # Evaluator evaluator = Evaluator(model) if args.evaluate: print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery) return # Criterion if args.arch == 'ResNet50_mgn_lr' or args.arch == 'ResNet101_mgn_lr' or args.arch == 'ResNet152_mgn_lr': criterion = MGN_loss(margin1=1.2, num_instances=4, alpha=1.0, gamma=1.0, theta=0.1, has_trip=True).cuda() elif args.arch == 'ResNet_reid_50' or args.arch == 'ResNet_reid_101' or args.arch == 'ResNet_reid_152': # criterion = XentropyLoss_SAC(theta=0.2,gamma=1).cuda() ranking_loss = nn.MarginRankingLoss(margin=args.margin).cuda() criterion = { 'XentropyLoss_SAC': XentropyLoss_SAC(theta=0.2,gamma=1).cuda(), \ 'trihard': TripletLoss(ranking_loss).cuda() } else: criterion = nn.CrossEntropyLoss().cuda() # Optimizer frozen_layerName = [ 'conv1', 'bn1', 'relu', 'maxpool', 'layer1', 'layer2', ] ##### Optimizer if args.frozen_sublayer: frozen_Source = None if hasattr(model.module, 'base'): frozen_Source = 'model.module.base.' elif hasattr(model.module, frozen_layerName[0]): frozen_Source = 'model.module.' else: raise RuntimeError( 'Not freeze layers but frozen_sublayer is True!') base_params_set = set() for subLayer in frozen_layerName: if hasattr(eval(frozen_Source[:-1]), subLayer): print('frozen layer: ', subLayer) single_module_param = eval(frozen_Source + subLayer + '.parameters()') # base_params.append(single_module_param) single_module_param_set = set(map(id, single_module_param)) base_params_set = base_params_set | single_module_param_set else: print("current model doesn't have ", subLayer) new_params = [ p for p in model.parameters() if id(p) not in base_params_set ] base_params = [ p for p in model.parameters() if id(p) in base_params_set ] param_groups = [{ 'params': base_params, 'lr_mult': 0.1 }, { 'params': new_params, 'lr_mult': 1.0 }] else: param_groups = model.parameters() optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Trainer if args.arch == 'ResNet50_mgn_lr' or args.arch == 'ResNet101_mgn_lr' or args.arch == 'ResNet152_mgn_lr': trainer = Trainer(model, criterion) elif args.arch == 'ResNet_reid_50' or args.arch == 'ResNet_reid_101' or args.arch == 'ResNet_reid_152': trainer = Trainer_SAC_Triplet( model, criterion, metric_loss_weight=args.metric_loss_weight) else: trainer = Cross_Trihard_Trainer( model, criterion, metric_loss_weight=args.metric_loss_weight) # Schedule learning rate print(args.step_epoch) scheduler = WarmupMultiStepLR(optimizer, args.step_epoch, gamma=args.gamma, warmup_factor=args.warm_up_factor, warmup_iters=args.warm_up_iter) # Start training for epoch in range(start_epoch + 1, args.epochs + 1): scheduler.step() trainer.train(epoch, train_loader, optimizer) if epoch % args.epoch_inter == 0 or epoch >= args.dense_evaluate: tmp_mAP, tmp_res = evaluator.evaluate(test_loader, dataset.query, dataset.gallery) if epoch >= args.start_save: if tmp_mAP > best: best = tmp_mAP flag = True else: flag = False # save_checkpoint({ # 'state_dict': model.module.state_dict(), # 'epoch': epoch, # 'best_map':tmp_mAP # }, flag, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch, 'best_map': tmp_mAP }, flag, fpath=osp.join(args.logs_dir, 'pass%d.pth.tar' % (epoch))) print('Test with best model:') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) model.module.load_state_dict(checkpoint['state_dict']) evaluator.evaluate(test_loader, dataset.query, dataset.gallery)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True logger = TensorLogger(osp.join(args.log_root, 'Tensorboard_logs', args.logs_dir)) # Redirect print to both console and log file logs_dir = osp.join(args.log_root, args.logs_dir) if not args.evaluate: sys.stdout = Logger(osp.join(logs_dir, 'log.txt')) # Create data loaders assert args.num_instances > 1, "num_instances should be greater than 1" assert args.source_batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) # Create source data_loader source_dataset, source_num_classes, source_train_loader\ , source_val_loader, source_test_loader = \ get_source_data(args.source_dataset, args.split, args.data_dir, args.height, args.width, args.source_batch_size, args.num_instances, args.workers, args.combine_trainval) # Create target data_loader target_dataset, target_num_classes, target_train_loader\ , target_val_loader, target_test_loader = \ get_target_data(args.target_dataset, args.split, args.data_dir, args.height, args.width, args.target_batch_size, args.workers, args.combine_trainval) print("lr:", args.lr) print("max_epoch:", args.epochs) # Create model # Hacking here to let the classifier be the last feature embedding layer # Net structure: avgpool -> FC(1024) -> FC(args.features) model = ResNet_recon(num_features=1024, dropout=args.dropout) if args.evaluate: model_evalu = nn.DataParallel(model).cuda() model = model.cuda(0) # For source triplet-loss trip_embedding = Trip_embedding(num_features=1024, num_diff_features=128, dropout=args.dropout).cuda(1) # For target reconstruction-loss recon_module = Reconstruct(num_features=1024).cuda(1) # Criterion criterion = ReconTripLoss().cuda(1) # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: model_path = load_checkpoint(args.resume) model.load_state_dict(model_path['state_dict']) # trip_embedding.load_state_dict(model_path['trip_em']) recon_module.load_state_dict(model_path['recon_dict']) start_epoch = model_path['epoch'] best_top1 = model_path['best_top1'] is_best = False top1 = best_top1 print("=> Start epoch {} best top1 {:.1%}" .format(start_epoch, best_top1)) # model = nn.DataParallel(model, device_ids=[0,1]).cuda(1) # model.cuda(0) # trip_embedding.cuda(1) # recon_module.cuda(1) # criterion.cuda(1) # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator if args.evaluate: evaluator = Evaluator(model_evalu) metric.train(model_evalu, source_train_loader) # print("Validation:") # evaluator.evaluate(val_loader, dataset.val, dataset.val, metric) print("Test:") if len(args.source_dataset) > 1: for dataset_name in args.source_dataset: print("{} source_test result:".format(dataset_name)) evaluator.evaluate(source_test_loader[dataset_name], source_dataset.query[dataset_name], source_dataset.gallery[dataset_name], metric) return else: print("source test result") evaluator.evaluate(source_test_loader, source_dataset.query, source_dataset.gallery, metric) print("target test result") evaluator.evaluate(target_test_loader, target_dataset.query, target_dataset.gallery, metric) return evaluator = Evaluator(model) # Optimizer optimizer = torch.optim.Adam([{'params': model.parameters()}, # {'params': trip_embedding.parameters()}, {'params': recon_module.parameters()}], lr=args.lr, weight_decay=args.weight_decay) # Trainer trainer = Transfer_Trainer(model, recon_module, trip_embedding, criterion) # Schedule learning rate def adjust_lr(epoch): lr = args.lr if epoch <= args.lr_change_epochs else \ args.lr * (0.001 ** ((epoch - args.lr_change_epochs)/ float(args.epochs-args.lr_change_epochs))) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # epoch_num = args.maxstep//(750//18) + 1 # Start training top1 = 0 is_best = True for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, source_train_loader, target_train_loader, optimizer, logger) #######Tensorboard-logs########## # for tag, value in model.named_parameters(): # tag = tag.replace('.', '/') # try: # logger.histo_summary(tag, to_np(value), epoch * len(source_train_loader) + 1) # logger.histo_summary(tag + '/grad', to_np(value.grad), epoch * len(source_train_loader) + 1) # except AttributeError, e: # pass if epoch % 25 == 0 and epoch != 0: top1 = evaluator.evaluate(source_test_loader, source_dataset.query, source_dataset.query) target_top1 = evaluator.evaluate(target_test_loader, target_dataset.query, target_dataset.query) print('target_top1 = {:5.1%}'.format(target_top1)) # top1 = evaluator.evaluate(test_loader, dataset.query, dataset.gallery) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint({ 'state_dict': model.state_dict(), 'recon_dict': recon_module.state_dict(), # 'trip_em': trip_embedding.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(epoch, top1, best_top1, ' *' if is_best else '')) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(logs_dir, 'checkpoint.pth.tar')) model.module.load_state_dict(checkpoint['state_dict']) metric.train(model, source_train_loader) print("Test:") if len(args.dataset) > 1: for dataset_name in args.dataset: print("{} test result:".format(dataset_name)) evaluator.evaluate(source_test_loader[dataset_name], source_dataset.query[dataset_name], source_dataset.gallery[dataset_name], metric) else: evaluator.evaluate(source_test_loader, source_dataset.query, source_dataset.gallery, metric)
def main(args): fix(args.seed) # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) print(args) # Create data loaders dataset, test_dataset, num_classes, source_train_loader, grid_query_loader, grid_gallery_loader,prid_query_loader, prid_gallery_loader,viper_query_loader, viper_gallery_loader, ilid_query_loader, ilid_gallery_loader = \ get_data(args.data_dir, args.height, args.width, args.batch_size, args.num_instance, args.re, args.workers) # Create model Encoder, Transfer, CamDis = models.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=num_classes) invNet = InvNet(args.features, num_classes, args.batch_size, beta=args.beta, knn=args.knn, alpha=args.alpha).cuda() # Load from checkpoint start_epoch = 0 if args.resume: checkpoint = load_checkpoint(args.resume) Encoder.load_state_dict(checkpoint['Encoder']) Transfer.load_state_dict(checkpoint['Transfer']) CamDis.load_state_dict(checkpoint['CamDis']) invNet.load_state_dict(checkpoint['InvNet']) start_epoch = checkpoint['epoch'] Encoder = Encoder.cuda() Transfer = Transfer.cuda() CamDis = CamDis.cuda() model = [Encoder, Transfer, CamDis] # Evaluator evaluator = Evaluator(model) if args.evaluate: # ----------------------------- v = evaluator.eval_viper(viper_query_loader, viper_gallery_loader, test_dataset.viper_query, test_dataset.viper_gallery, args.output_feature, seed=57) p = evaluator.eval_prid(prid_query_loader, prid_gallery_loader, test_dataset.prid_query, test_dataset.prid_gallery, args.output_feature, seed=40) g = evaluator.eval_grid(grid_query_loader, grid_gallery_loader, test_dataset.grid_query, test_dataset.grid_gallery, args.output_feature, seed=35) l = evaluator.eval_ilids(ilid_query_loader, test_dataset.ilid_query, args.output_feature, seed=24) # ----------------------------- criterion = [] criterion.append(nn.CrossEntropyLoss().cuda()) criterion.append(TripletLoss(margin=args.margin)) # Optimizer base_param_ids = set(map(id, Encoder.base.parameters())) new_params = [p for p in Encoder.parameters() if id(p) not in base_param_ids] param_groups = [ {'params': Encoder.base.parameters(), 'lr_mult': 0.1}, {'params': new_params, 'lr_mult': 1.0}] optimizer_Encoder = torch.optim.SGD(param_groups, lr=args.lr, momentum=0.9, weight_decay=5e-4, nesterov=True) # ==== base_param_ids = set(map(id, Transfer.base.parameters())) new_params = [p for p in Transfer.parameters() if id(p) not in base_param_ids] param_groups = [ {'params': Transfer.base.parameters(), 'lr_mult': 0.1}, {'params': new_params, 'lr_mult': 1.0}] optimizer_Transfer = torch.optim.SGD(param_groups, lr=args.lr, momentum=0.9, weight_decay=5e-4, nesterov=True) # ==== param_groups = [ {'params':CamDis.parameters(), 'lr_mult':1.0}, ] optimizer_Cam = torch.optim.SGD(param_groups, lr=args.lr,momentum=0.9, weight_decay=5e-4, nesterov=True) optimizer = [optimizer_Encoder, optimizer_Transfer, optimizer_Cam] # Trainer trainer = Trainer(model, criterion, InvNet=invNet) # Schedule learning rate def adjust_lr(epoch): step_size = 40 lr = args.lr * (0.1 ** ((epoch) // step_size)) for g in optimizer_Encoder.param_groups: g['lr'] = lr * g.get('lr_mult', 1) for g in optimizer_Transfer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) for g in optimizer_Cam.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, source_train_loader, optimizer, args.tri_weight, args.adv_weight, args.mem_weight) save_checkpoint({ 'Encoder': Encoder.state_dict(), 'Transfer': Transfer.state_dict(), 'CamDis': CamDis.state_dict(), 'InvNet': invNet.state_dict(), 'epoch': epoch + 1, }, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) evaluator = Evaluator(model) print('\n * Finished epoch {:3d} \n'. format(epoch)) # Final test print('Test with best model:') evaluator = Evaluator(model) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, args.output_feature, args.rerank)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders assert args.num_instances > 1, "num_instances should be greater than 1" assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) dataset, num_classes, train_loader, train_loader_head, train_loader_upper, train_loader_lower,\ val_loader, val_loader_head, val_loader_upper, val_loader_lower,\ test_loader, test_loader_head, test_loader_upper, test_loader_lower= \ get_data(args.dataset, args.split, args.data_dir, args.height,args.width, args.batch_size,args.num_instances,args.workers,args.combine_trainval) # Create model # Hacking here to let the classifier be the last feature embedding layer # Net structure: avgpool -> FC(1024) -> FC(args.features) model = models.create(args.arch, num_features=1024, dropout=args.dropout, num_classes=args.features) model_head = models.create(args.arch, num_features=1024, dropout=args.dropout, num_classes=args.features) model_upper = models.create(args.arch, num_features=1024, dropout=args.dropout, num_classes=args.features) model_lower = models.create(args.arch, num_features=1024, dropout=args.dropout, num_classes=args.features) # Load from checkpoint start_epoch = best_top1 = 0 # if args.resume: # checkpoint = load_checkpoint(args.resume) # model.load_state_dict(checkpoint['state_dict']) # start_epoch = checkpoint['epoch'] # best_top1 = checkpoint['best_top1'] # print("=> Start epoch {} best top1 {:.1%}" # .format(start_epoch, best_top1)) model = nn.DataParallel(model).cuda() model_head = nn.DataParallel(model_head).cuda() model_upper = nn.DataParallel(model_upper).cuda() model_lower = nn.DataParallel(model_lower).cuda() # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model, model_head, model_upper, model_lower) #if args.evaluate: # metric.train(model, train_loader) # print("Validation:") # evaluator.evaluate(val_loader, dataset.val, dataset.val, metric) # print("Test:") # evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric) # return # Criterion criterion = TripletLoss(margin=args.margin).cuda() criterion_head = TripletLoss(margin=args.margin).cuda() criterion_upper = TripletLoss(margin=args.margin).cuda() criterion_lower = TripletLoss(margin=args.margin).cuda() # Optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) optimizer_head = torch.optim.Adam(model_head.parameters(), lr=args.lr, weight_decay=args.weight_decay) optimizer_upper = torch.optim.Adam(model_upper.parameters(), lr=args.lr, weight_decay=args.weight_decay) optimizer_lower = torch.optim.Adam(model_lower.parameters(), lr=args.lr, weight_decay=args.weight_decay) # Trainer trainer = Trainer(model, criterion) trainer_head = Trainer(model_head, criterion_head) trainer_upper = Trainer(model_upper, criterion_upper) trainer_lower = Trainer(model_lower, criterion_lower) # Schedule learning rate def adjust_lr(epoch): lr = args.lr if epoch <= 100 else \ args.lr * (0.001 ** ((epoch - 100) / 50.0)) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) def adjust_lr_head(epoch): lr = args.lr if epoch <= 100 else \ args.lr * (0.001 ** ((epoch - 100) / 50.0)) for g in optimizer_head.param_groups: g['lr'] = lr * g.get('lr_mult', 1) def adjust_lr_upper(epoch): lr = args.lr if epoch <= 100 else \ args.lr * (0.001 ** ((epoch - 100) / 50.0)) for g in optimizer_upper.param_groups: g['lr'] = lr * g.get('lr_mult', 1) def adjust_lr_lower(epoch): lr = args.lr if epoch <= 100 else \ args.lr * (0.001 ** ((epoch - 100) / 50.0)) for g in optimizer_lower.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) adjust_lr_head(epoch) adjust_lr_upper(epoch) adjust_lr_lower(epoch) trainer.train(epoch, train_loader, optimizer) trainer_head.train(epoch, train_loader_head, optimizer_head) trainer_upper.train(epoch, train_loader_upper, optimizer_upper) trainer_lower.train(epoch, train_loader_lower, optimizer_lower) if epoch < args.start_save: continue top1 = evaluator.evaluate(val_loader, val_loader_head, val_loader_upper, val_loader_lower, dataset.val, dataset.val) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'), opath='model_best.pth.tar') save_checkpoint( { 'state_dict': model_head.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint_head.pth.tar'), opath='model_head_best.pth.tar') save_checkpoint( { 'state_dict': model_upper.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint_upper.pth.tar'), opath='model_upper_best.pth.tar') save_checkpoint( { 'state_dict': model_lower.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint_lower.pth.tar'), opath='model_lower_best.pth.tar') print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(epoch, top1, best_top1, ' *' if is_best else '')) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) checkpoint_head = load_checkpoint( osp.join(args.logs_dir, 'model_head_best.pth.tar')) checkpoint_upper = load_checkpoint( osp.join(args.logs_dir, 'model_upper_best.pth.tar')) checkpoint_lower = load_checkpoint( osp.join(args.logs_dir, 'model_lower_best.pth.tar')) model.module.load_state_dict(checkpoint['state_dict']) model_head.module.load_state_dict(checkpoint_head['state_dict']) model_upper.module.load_state_dict(checkpoint_upper['state_dict']) model_lower.module.load_state_dict(checkpoint_lower['state_dict']) metric.train(model, train_loader) metric.train(model_head, train_loader_head) metric.train(model_upper, train_loader_upper) metric.train(model_lower, train_loader_lower) evaluator.evaluate(test_loader, test_loader_head, test_loader_upper, test_loader_lower, dataset.query, dataset.gallery, metric)
def main(args): # For fast training. cudnn.benchmark = True device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) print('log_dir=', args.logs_dir) # Print logs print(args) # Create data loaders dataset, num_classes, source_train_loader, target_train_loader, \ query_loader, gallery_loader = get_data(args.data_dir, args.source, args.target, args.height, args.width, args.batch_size, args.camstyle_type, args.re, args.workers) # Create model model = models.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=num_classes) # Invariance learning model num_tgt = len(dataset.target_train) model_inv = InvNet(args.features, num_tgt, beta=args.inv_beta, knn=args.knn, alpha=args.inv_alpha) # Load from checkpoint start_epoch = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) model_inv.load_state_dict(checkpoint['state_dict_inv']) start_epoch = checkpoint['epoch'] print("=> Start epoch {} ".format(start_epoch)) # Set model model = nn.DataParallel(model).to(device) model_inv = model_inv.to(device) # Evaluator evaluator = Evaluator(model) if args.evaluate: print("Test:") query, retrieveds = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, args.output_feature) for i in range(len(query)): Image.open("data/" + args.target + "/query/" + query[i]).save( "q_%d.jpg" % i) for j in range(len(retrieveds[i])): Image.open("data/" + args.target + "/bounding_box_test/" + retrieveds[i][j]).save("r_%d_%d.jpg" % (i, j)) return # Optimizer base_param_ids = set(map(id, model.module.base.parameters())) base_params_need_for_grad = filter(lambda p: p.requires_grad, model.module.base.parameters()) new_params = [p for p in model.parameters() if id(p) not in base_param_ids] param_groups = [{ 'params': base_params_need_for_grad, 'lr_mult': 0.1 }, { 'params': new_params, 'lr_mult': 1.0 }] optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Trainer trainer = Trainer(model, model_inv, lmd=args.lmd, include_mmd=args.include_mmd, include_coral=args.include_coral, lmd_ext=args.lmd_ext) # Schedule learning rate def adjust_lr(epoch): step_size = args.epochs_decay lr = args.lr * (0.1**(epoch // step_size)) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, source_train_loader, target_train_loader, optimizer) save_checkpoint( { 'state_dict': model.module.state_dict(), 'state_dict_inv': model_inv.state_dict(), 'epoch': epoch + 1, }, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} \n'.format(epoch)) # Final test print('Test with best model:') evaluator = Evaluator(model) query, retrieveds = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, args.output_feature) for i in range(len(query)): Image.open("data/" + args.target + "/query/" + query[i]).save( "q_%d.jpg" % i) for j in range(len(retrieveds[i])): Image.open("data/" + args.target + "/bounding_box_test/" + retrieveds[i][j]).save("r_%d_%d.jpg" % (i, j))
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) random.seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file date_str = '{}'.format(datetime.datetime.today().strftime('%Y-%m-%d_%H-%M-%S')) if (not args.evaluate) and args.log: sys.stdout = Logger(osp.join(args.logs_dir, 'log_{}.txt'.format(date_str))) # save opts with open(osp.join(args.logs_dir, 'args_{}.json'.format(date_str)), 'w') as fp: json.dump(vars(args), fp, indent=1) assert args.real or args.synthetic, "At least one dataset should be used" # Create data loaders print (args.real, args.synthetic) dataset, num_classes, train_loader, query_loader, gallery_loader, camstyle_loader = \ get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.num_workers, args.combine_trainval, args.crop, args.tracking_icams, args.tracking_fps, args.real, args.synthetic, args.re, 0, args.camstyle) # Create model model = models.create('ide', num_features=args.features, norm=args.norm, dropout=args.dropout, num_classes=num_classes, last_stride=args.last_stride, output_feature=args.output_feature, arch=args.arch) # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: if args.evaluate: model, start_epoch, best_top1 = checkpoint_loader(model, args.resume, eval_only=True) else: model, start_epoch, best_top1 = checkpoint_loader(model, args.resume) print("=> Start epoch {} best top1_eval {:.1%}".format(start_epoch, best_top1)) model = nn.DataParallel(model).cuda() # Evaluator evaluator = Evaluator(model) if args.evaluate: print("Test:") evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, eval_only=True) return if args.evaluate_VID: print("Test on VID dataset:") mAP_list = [] cmc1_list = [] cmc5_list = [] cmc10_list = [] for i in range(10): dataset, num_classes, train_loader, query_loader, gallery_loader, camstyle_loader = \ get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.num_workers, args.combine_trainval, args.crop, args.tracking_icams, args.tracking_fps, args.re, 0, args.camstyle) mAP, cmc1, cmc5 , cmc10 = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, eval_only=True) mAP_list.append(mAP) cmc1_list.append(cmc1) cmc5_list.append(cmc5) cmc10_list.append(cmc10) print('Final VID test [mAP: {:5.2%}], [cmc1: {:5.2%}], [cmc5: {:5.2%}], [cmc10: {:5.2%}]' .format(np.mean(mAP_list), np.mean(cmc1_list), np.mean(cmc5_list), np.mean(cmc10_list))) return # Criterion criterion = nn.CrossEntropyLoss().cuda() if not args.LSR else LSR_loss().cuda() if args.train or args.finetune: # Optimizer if hasattr(model.module, 'base'): # low learning_rate the base network (aka. ResNet-50) base_param_ids = set(map(id, model.module.base.parameters())) new_params = [p for p in model.parameters() if id(p) not in base_param_ids] param_groups = [{'params': model.module.base.parameters(), 'lr_mult': 0.1}, {'params': new_params, 'lr_mult': 1.0}] else: param_groups = model.parameters() optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Trainer if args.camstyle == 0: trainer = Trainer(model, criterion) else: trainer = CamStyleTrainer(model, criterion, camstyle_loader) # Schedule learning rate def adjust_lr(epoch): step_size = args.step_size lr = args.lr * (0.1 ** (epoch // step_size)) if args.finetune: if args.dataset == "veri": lr = lr / 50 elif args.dataset == "vihicle_id": lr = lr / 10 for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Draw Curve epoch_s = [] loss_s = [] prec_s = [] # Start training for epoch in range(start_epoch, args.epochs): t0 = time.time() adjust_lr(epoch) # train_loss, train_prec = 0, 0 train_loss, train_prec = trainer.train(epoch, train_loader, optimizer, fix_bn=args.fix_bn) if epoch < args.start_save: continue if epoch % 20 == 0 or args.finetune: evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, eval_only=True) top1_eval = 50 is_best = top1_eval >= best_top1 best_top1 = max(top1_eval, best_top1) save_checkpoint({ 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint_{}.pth.tar'.format(date_str))) epoch_s.append(epoch) loss_s.append(train_loss) prec_s.append(train_prec) draw_curve(os.path.join(args.logs_dir, 'train_{}.jpg'.format(date_str)), epoch_s, loss_s, prec_s) t1 = time.time() t_epoch = t1 - t0 print('\n * Finished epoch {:3d} top1_eval: {:5.1%} best_eval: {:5.1%} \n'. format(epoch, top1_eval, best_top1, ' *' if is_best else '')) print('*************** Epoch takes time: {:^10.2f} *********************\n'.format(t_epoch)) pass # Final test print('Test with best model:') model, start_epoch, best_top1 = checkpoint_loader(model, osp.join(args.logs_dir, 'model_best.pth.tar'), eval_only=True) print("=> Start epoch {} best top1 {:.1%}".format(start_epoch, best_top1)) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, eval_only=True)
def main(args): cudnn.deterministic = False cudnn.benchmark = True exp_database_dir = osp.join(args.exp_dir, string.capwords(args.dataset)) output_dir = osp.join(exp_database_dir, args.method, args.sub_method) log_file = osp.join(output_dir, 'log.txt') # Redirect print to both console and log file sys.stdout = Logger(log_file) # Create model ibn_type = args.ibn if ibn_type == 'none': ibn_type = None model = resmap.create(args.arch, ibn_type=ibn_type, final_layer=args.final_layer, neck=args.neck).cuda() num_features = model.num_features # print(model) # print('\n') feamap_factor = {'layer2': 8, 'layer3': 16, 'layer4': 32} hei = args.height // feamap_factor[args.final_layer] wid = args.width // feamap_factor[args.final_layer] matcher = QAConv(num_features, hei, wid).cuda() for arg in sys.argv: print('%s ' % arg, end='') print('\n') # Criterion criterion = PairwiseMatchingLoss(matcher).cuda() # Optimizer base_param_ids = set(map(id, model.base.parameters())) new_params = [p for p in model.parameters() if id(p) not in base_param_ids] param_groups = [{ 'params': model.base.parameters(), 'lr': 0.1 * args.lr }, { 'params': new_params, 'lr': args.lr }, { 'params': matcher.parameters(), 'lr': args.lr }] optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=0.9, weight_decay=5e-4, nesterov=True) # Load from checkpoint start_epoch = 0 if args.resume or args.evaluate: print('Loading checkpoint...') if args.resume and (args.resume != 'ori'): checkpoint = load_checkpoint(args.resume) else: checkpoint = load_checkpoint( osp.join(output_dir, 'checkpoint.pth.tar')) model.load_state_dict(checkpoint['model']) criterion.load_state_dict(checkpoint['criterion']) optimizer.load_state_dict(checkpoint['optim']) start_epoch = checkpoint['epoch'] print("=> Start epoch {} ".format(start_epoch)) model = nn.DataParallel(model).cuda() # Create data loaders # Warning: this training data loader cannot be used elsewhere other than a continueous training, otherwise the # switch between the PK sampler and GS sampler will be incorrect! save_path = None if args.gs_save: save_path = output_dir dataset, num_classes, train_loader, _, _ = get_data( args.dataset, args.data_dir, model, matcher, start_epoch - 1, save_path, args) # Decay LR by a factor of 0.1 every step_size epochs lr_scheduler = StepLR(optimizer, step_size=args.step_size, gamma=0.1, last_epoch=start_epoch - 1) if not args.evaluate: # Trainer trainer = Trainer(model, criterion, args.clip_value) t0 = time.time() # Start training for epoch in range(start_epoch, args.epochs): loss, acc = trainer.train(epoch, train_loader, optimizer) lr = list(map(lambda group: group['lr'], optimizer.param_groups)) lr_scheduler.step() train_time = time.time() - t0 print( '* Finished epoch %d at lr=[%g, %g, %g]. Loss: %.3f. Acc: %.2f%%. Training time: %.0f seconds. \n' % (epoch + 1, lr[0], lr[1], lr[2], loss, acc * 100, train_time)) save_checkpoint( { 'model': model.module.state_dict(), 'criterion': criterion.state_dict(), 'optim': optimizer.state_dict(), 'epoch': epoch + 1, }, fpath=osp.join(output_dir, 'checkpoint.pth.tar')) # Final test print('Evaluate the learned model:') t0 = time.time() # Evaluator evaluator = Evaluator(model) avg_rank1 = 0 avg_mAP = 0 num_testsets = 0 results = {} test_names = args.testset.strip().split(',') for test_name in test_names: if test_name not in datasets.names(): print('Unknown dataset: %s.' % test_name) continue testset, test_query_loader, test_gallery_loader = \ get_test_data(test_name, args.data_dir, args.height, args.width, args.workers, args.test_fea_batch) if not args.do_tlift: testset.has_time_info = False test_rank1, test_mAP, test_rank1_rerank, test_mAP_rerank, test_rank1_tlift, test_mAP_tlift, test_dist, \ test_dist_rerank, test_dist_tlift, pre_tlift_dict = \ evaluator.evaluate(matcher, testset, test_query_loader, test_gallery_loader, args.test_gal_batch, args.test_prob_batch, args.tau, args.sigma, args.K, args.alpha) results[test_name] = [test_rank1, test_mAP] if test_name != args.dataset: avg_rank1 += test_rank1 avg_mAP += test_mAP num_testsets += 1 if testset.has_time_info: print( ' %s: rank1=%.1f, mAP=%.1f, rank1_rerank=%.1f, mAP_rerank=%.1f,' ' rank1_rerank_tlift=%.1f, mAP_rerank_tlift=%.1f.\n' % (test_name, test_rank1 * 100, test_mAP * 100, test_rank1_rerank * 100, test_mAP_rerank * 100, test_rank1_tlift * 100, test_mAP_tlift * 100)) else: print(' %s: rank1=%.1f, mAP=%.1f.\n' % (test_name, test_rank1 * 100, test_mAP * 100)) result_file = osp.join(exp_database_dir, args.method, test_name + '_results.txt') with open(result_file, 'a') as f: f.write('%s/%s:\n' % (args.method, args.sub_method)) if testset.has_time_info: f.write( '\t%s: rank1=%.1f, mAP=%.1f, rank1_rerank=%.1f, mAP_rerank=%.1f, rank1_rerank_tlift=%.1f, ' 'mAP_rerank_tlift=%.1f.\n\n' % (test_name, test_rank1 * 100, test_mAP * 100, test_rank1_rerank * 100, test_mAP_rerank * 100, test_rank1_tlift * 100, test_mAP_tlift * 100)) else: f.write('\t%s: rank1=%.1f, mAP=%.1f.\n\n' % (test_name, test_rank1 * 100, test_mAP * 100)) if args.save_score: test_gal_list = np.array( [fname for fname, _, _, _ in testset.gallery], dtype=np.object) test_prob_list = np.array( [fname for fname, _, _, _ in testset.query], dtype=np.object) test_gal_ids = [pid for _, pid, _, _ in testset.gallery] test_prob_ids = [pid for _, pid, _, _ in testset.query] test_gal_cams = [c for _, _, c, _ in testset.gallery] test_prob_cams = [c for _, _, c, _ in testset.query] test_score_file = osp.join(exp_database_dir, args.method, args.sub_method, '%s_score.mat' % test_name) sio.savemat(test_score_file, { 'score': 1. - test_dist, 'score_rerank': 1. - test_dist_rerank, 'score_tlift': 1. - test_dist_tlift, 'gal_time': pre_tlift_dict['gal_time'], 'prob_time': pre_tlift_dict['prob_time'], 'gal_list': test_gal_list, 'prob_list': test_prob_list, 'gal_ids': test_gal_ids, 'prob_ids': test_prob_ids, 'gal_cams': test_gal_cams, 'prob_cams': test_prob_cams }, oned_as='column', do_compression=True) test_time = time.time() - t0 avg_rank1 /= num_testsets avg_mAP /= num_testsets for key in results.keys(): print('%s: rank1=%.1f%%, mAP=%.1f%%.' % (key, results[key][0] * 100, results[key][1] * 100)) print('Average: rank1=%.2f%%, mAP=%.2f%%.\n\n' % (avg_rank1 * 100, avg_mAP * 100)) result_file = osp.join(exp_database_dir, args.method, args.sub_method[:-5] + '_avg_results.txt') with open(result_file, 'a') as f: f.write('%s/%s:\n' % (args.method, args.sub_method)) if not args.evaluate: f.write('\t Loss: %.3f, acc: %.2f%%. ' % (loss, acc * 100)) f.write("Train: %.0fs. " % train_time) f.write("Test: %.0fs. " % test_time) f.write('Rank1: %.2f%%, mAP: %.2f%%.\n' % (avg_rank1 * 100, avg_mAP * 100)) for key in results.keys(): f.write('\t %s: Rank1: %.1f%%, mAP: %.1f%%.\n' % (key, results[key][0] * 100, results[key][1] * 100)) f.write('\n') if not args.evaluate: print('Finished training at epoch %d, loss = %.3f, acc = %.2f%%.\n' % (epoch + 1, loss, acc * 100)) print( "Total training time: %.3f sec. Average training time per epoch: %.3f sec." % (train_time, train_time / (args.epochs - start_epoch + 1))) print("Total testing time: %.3f sec.\n" % test_time) for arg in sys.argv: print('%s ' % arg, end='') print('\n')
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Create data loaders assert args.num_instances > 1, "num_instances should be greater than 1" assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) # get source data src_dataset, src_extfeat_loader = \ get_source_data(args.src_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # get target data tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \ get_data(args.tgt_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # Create model # Hacking here to let the classifier be the number of source ids if args.src_dataset == 'dukemtmc': model = models.create(args.arch, num_classes=632, pretrained=False) coModel = models.create(args.arch, num_classes=632, pretrained=False) elif args.src_dataset == 'market1501': model = models.create(args.arch, num_classes=676, pretrained=False) coModel = models.create(args.arch, num_classes=676, pretrained=False) elif args.src_dataset == 'msmt17': model = models.create(args.arch, num_classes=1041, pretrained=False) coModel = models.create(args.arch, num_classes=1041, pretrained=False) elif args.src_dataset == 'cuhk03': model = models.create(args.arch, num_classes=1230, pretrained=False) coModel = models.create(args.arch, num_classes=1230, pretrained=False) else: raise RuntimeError('Please specify the number of classes (ids) of the network.') # Load from checkpoint if args.resume: print('Resuming checkpoints from finetuned model on another dataset...\n') checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict'], strict=False) coModel.load_state_dict(checkpoint['state_dict'], strict=False) else: raise RuntimeWarning('Not using a pre-trained model.') model = nn.DataParallel(model).cuda() coModel = nn.DataParallel(coModel).cuda() evaluator = Evaluator(model, print_freq=args.print_freq) # evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) # if args.evaluate: return # Criterion criterion = [ TripletLoss(args.margin, args.num_instances, isAvg=False, use_semi=False).cuda(), TripletLoss(args.margin, args.num_instances, isAvg=False, use_semi=False).cuda(), ] # Optimizer optimizer = torch.optim.Adam( model.parameters(), lr=args.lr ) coOptimizer = torch.optim.Adam( coModel.parameters(), lr=args.lr ) optims = [optimizer, coOptimizer] # training stage transformer on input images normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transformer = T.Compose([ T.Resize((args.height, args.width)), T.RandomHorizontalFlip(), T.ToTensor(), normalizer, T.RandomErasing(probability=0.5, sh=0.2, r1=0.3) ]) # # Start training for iter_n in range(args.iteration): if args.lambda_value == 0: source_features = 0 else: # get source datas' feature source_features, _ = extract_features(model, src_extfeat_loader, print_freq=args.print_freq, numStripe=None) # synchronization feature order with src_dataset.train source_features = torch.cat([source_features[f].unsqueeze(0) for f, _, _ in src_dataset.train], 0) # extract training images' features print('Iteration {}: Extracting Target Dataset Features...'.format(iter_n + 1)) target_features, _ = extract_features(model, tgt_extfeat_loader, print_freq=args.print_freq, numStripe=None) # synchronization feature order with dataset.train target_features = torch.cat([target_features[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval], 0) # calculate distance and rerank result print('Calculating feature distances...') target_features = target_features.numpy() rerank_dist = re_ranking(source_features, target_features, lambda_value=args.lambda_value) if iter_n == 0: # DBSCAN cluster tri_mat = np.triu(rerank_dist, 1) # tri_mat.dim=2 tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1 tri_mat = np.sort(tri_mat, axis=None) top_num = np.round(args.rho * tri_mat.size).astype(int) eps = tri_mat[:top_num].mean() print('eps in cluster: {:.3f}'.format(eps)) cluster = DBSCAN(eps=eps, min_samples=4, metric='precomputed', n_jobs=8) # select & cluster images as training set of this epochs print('Clustering and labeling...') labels = cluster.fit_predict(rerank_dist) num_ids = len(set(labels)) - 1 print('Iteration {} have {} training ids'.format(iter_n + 1, num_ids)) # generate new dataset new_dataset, unknown_dataset = [], [] # assign label for target ones unknownLab = labelNoise(torch.from_numpy(target_features), torch.from_numpy(labels)) # unknownFeats = target_features[labels==-1,:] unCounter, index = 0, 0 from collections import defaultdict realIDs, fakeIDs = defaultdict(list), [] for (fname, realPID, cam), label in zip(tgt_dataset.trainval, labels): if label == -1: unknown_dataset.append((fname, int(unknownLab[unCounter]), cam)) # unknown data fakeIDs.append(int(unknownLab[unCounter])) realIDs[realPID].append(index) unCounter += 1 index += 1 continue # dont need to change codes in trainer.py _parsing_input function and sampler function after add 0 new_dataset.append((fname, label, cam)) fakeIDs.append(label) realIDs[realPID].append(index) index += 1 print('Iteration {} have {} training images'.format(iter_n + 1, len(new_dataset))) precision, recall, fscore = calScores(realIDs, np.asarray(fakeIDs)) # fakeIDs does not contain -1 print('precision:{}, recall:{}, fscore: {}'.format(100 * precision, 100 * recall, fscore)) train_loader = DataLoader( Preprocessor(new_dataset, root=tgt_dataset.images_dir, transform=train_transformer), batch_size=args.batch_size, num_workers=4, sampler=RandomIdentitySampler(new_dataset, args.num_instances), pin_memory=True, drop_last=True ) # hard samples # noiseImgs = [name[1] for name in unknown_dataset] # saveAll(noiseImgs, tgt_dataset.images_dir, 'noiseImg') # import ipdb; ipdb.set_trace() unLoader = DataLoader( Preprocessor(unknown_dataset, root=tgt_dataset.images_dir, transform=train_transformer), batch_size=args.batch_size, num_workers=4, sampler=RandomIdentitySampler(unknown_dataset, args.num_instances), pin_memory=True, drop_last=True ) # train model with new generated dataset trainer = RCoTeaching( model, coModel, train_loader, unLoader, criterion, optims ) # Start training for epoch in range(args.epochs): trainer.train(epoch, remRate=0.2 + (0.8 / args.iteration) * (1 + iter_n)) # test only rank_score = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) # print('co-model:\n') # rank_score = evaluatorB.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) # Evaluate rank_score = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) save_checkpoint({ 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': rank_score.market1501[0], }, True, fpath=osp.join(args.logs_dir, 'RCT.pth')) return rank_score.map, rank_score.market1501[0]
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) dataset, num_classes, train_loader, query_loader, gallery_loader = \ get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers, ) # Create model model = models.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=num_classes) # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model_dict = model.state_dict() checkpoint_load = { k: v for k, v in (checkpoint['state_dict']).items() if k in model_dict } model_dict.update(checkpoint_load) model.load_state_dict(model_dict) # model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] best_top1 = checkpoint['best_top1'] print("=> Start epoch {} best top1 {:.1%}".format( start_epoch, best_top1)) model = nn.DataParallel(model).cuda() # Evaluator evaluator = Evaluator(model) if args.evaluate: print("Test:") evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) return # Criterion criterion = nn.CrossEntropyLoss().cuda() # Optimizer if hasattr(model.module, 'base'): base_param_ids = set(map(id, model.module.base.parameters())) new_params = [ p for p in model.parameters() if id(p) not in base_param_ids ] param_groups = [{ 'params': model.module.base.parameters(), 'lr_mult': 0.1 }, { 'params': new_params, 'lr_mult': 1.0 }] else: param_groups = model.parameters() optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # optimizer = torch.optim.Adam(param_groups,lr=args.lr) # Trainer trainer = Trainer(model, criterion, 0, 0, SMLoss_mode=0) # Schedule learning rate def adjust_lr(epoch): step_size = 60 if args.arch == 'inception' else args.step_size lr = args.lr * (0.1**(epoch // step_size)) # if epoch>70: # lr = 0.01 for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, train_loader, optimizer) is_best = True save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'checkpoint.pth.tar')) model.module.load_state_dict(checkpoint['state_dict']) results = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) # save the parameters and results with open('%s/opts.yaml' % args.logs_dir, 'w') as fp: yaml.dump(vars(args), fp, default_flow_style=False) txtName = args.logs_dir + "results.txt" file = open(txtName, 'w') for key in results: file.write(key + ': ' + str(results[key]) + '\n')
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders assert args.num_instances > 1, "num_instances should be greater than 1" assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (256, 128) dataset, num_classes, train_loader, val_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.height, args.width, args.batch_size, args.num_instances, args.workers, args.combine_trainval) # Create model # Hacking here to let the classifier be the last feature embedding layer # Net structure: avgpool -> FC(1024) -> FC(args.features) model = ResNet(args.depth, num_features=args.features, dropout=args.dropout, num_classes=num_classes) model = nn.DataParallel(model).cuda() # Load from checkpoint start_epoch = best_map = 0 if args.if_resume: print(args.resume) checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) prior_best_map = checkpoint['best_top1'] print("=> Start epoch {} best top1 {:.1%}" .format(start_epoch, prior_best_map)) # model = nn.DataParallel(model).cuda() # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model) if args.evaluate: metric.train(model, train_loader) print("Validation:") evaluator.evaluate(val_loader, dataset.val, dataset.val, metric) print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric) return # Criterion alpha= args.alpha beta = args.beta gamma = args.gamma criterion = TripletLoss_biu(margin = args.margin, num_instances=args.num_instances, alpha = alpha, beta =beta , gamma =gamma).cuda() # Optimizer if args.optimizer == 'sgd': # base_param_ids = set(map(id, model.module.base.parameters())) # new_params = [p for p in model.parameters() if id(p) not in base_param_ids] # param_groups = [ # {'params': model.module.base.parameters(), 'lr_mult': 0.1}, # {'params': new_params, 'lr_mult': 1.0}] param_groups = model.parameters() optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) else : optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # Trainer trainer = Trainer(model, criterion) # Schedule learning rate def adjust_lr(epoch): if args.optimizer == 'sgd': lr = args.lr * (0.1 ** (epoch // 40)) else : lr = args.lr if epoch <= 80 else \ args.lr * (0.1 ** ((epoch - 100) / 60.0)) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, train_loader, optimizer) if epoch % 3 ==0: metric.train(model,train_loader) top_map = evaluator.evaluate(test_loader, dataset.query, dataset.gallery) is_best = top_map > prior_best_map prior_best_map = max(top_map, prior_best_map) save_checkpoint({ 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_map': top_map, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) # print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. # format(epoch, top1, best_top1, ' *' if is_best else '')) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) model.module.load_state_dict(checkpoint['state_dict']) metric.train(model, train_loader) evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir + '/log')) if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else (256, 128) dataset, num_classes, train_loader, trainvallabel, val_loader, query_loader, gallery_loader = \ get_data(args.dataset, args.split, args.data_dir, args.height, args.width, args.batch_size, args.num_instances, args.workers, args.combine_trainval, args.flip_prob, args.padding, args.re_prob) model = models.create(args.arch, num_classes=num_classes, num_features=args.features, attention_mode=args.att_mode) print(model) model = model.cuda() evaluator = Evaluator(model) metric = DistanceMetric(algorithm=args.dist_metric) start_epoch = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['model']) start_epoch = checkpoint['epoch'] print("=> Start epoch {}".format(start_epoch)) if args.evaluate: metric.train(model, train_loader) evaluator.evaluate(dataset.query, dataset.gallery) exit() current_margin = args.margin criterion_z = nn.CrossEntropyLoss().cuda() # criterion_z = CrossEntropyLabelSmooth(num_classes= num_classes, epsilon=0.5).cuda() # criterion_I = TripletLoss(margin= current_margin).cuda() # criterion_D = nn.CrossEntropyLoss().cuda() print(args) if args.arch == 'ide': ignored_params = list(map(id, model.model.fc.parameters())) + list( map(id, model.classifier.parameters())) else: ignored_params = list(map(id, model.classifier.parameters())) + list( map(id, model.attention_module.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) if args.use_adam: # optimizer_ft = torch.optim.Adam(model.parameters(),lr=args.lr, weight_decay=5e-4) optimizer_ft = torch.optim.Adam([ { 'params': filter(lambda p: p.requires_grad, base_params), 'lr': args.lr }, { 'params': model.classifier.parameters(), 'lr': args.lr }, { 'params': model.attention_module.parameters(), 'lr': args.lr }, ], weight_decay=5e-4) # optimizer_discriminator = torch.optim.Adam([ # {'params': model_discriminator.model.parameters(), 'lr': args.lr}, # ], # weight_decay=5e-4) else: optimizer_ft = torch.optim.Adam([ { 'params': filter(lambda p: p.requires_grad, base_params), 'lr': args.lr }, { 'params': model.classifier.parameters(), 'lr': args.lr }, { 'params': model.attention_module.parameters(), 'lr': args.lr }, ], weight_decay=5e-4) # optimizer_ft = torch.optim.SGD([ # {'params': filter(lambda p: p.requires_grad,base_params), 'lr': args.lr}, # {'params': model.classifier.parameters(), 'lr': args.lr}, # {'params': model.att_classifier.parameters(), 'lr': args.lr}, # {'params': model.projector1.parameters(), 'lr': args.lr}, # {'params': model.projector2.parameters(), 'lr': args.lr}, # {'params': model.projector3.parameters(), 'lr': args.lr}, # {'params': model.projector4.parameters(), 'lr': args.lr}, # {'params': model.att1.parameters(), 'lr': args.lr}, # {'params': model.att2.parameters(), 'lr': args.lr}, # {'params': model.att3.parameters(), 'lr': args.lr}, # {'params': model.att4.parameters(), 'lr': args.lr}, # ], # momentum=0.9, # weight_decay=5e-4, # nesterov=True) # optimizer_discriminator = torch.optim.SGD([ # {'params': model_discriminator.model.parameters(), 'lr': args.lr}, # {'params': model_discriminator.classifier.parameters(), 'lr': args.lr}, # ], # momentum=0.9, # weight_decay=5e-4, # nesterov=True) scheduler = WarmupMultiStepLR(optimizer_ft, args.mile_stone, args.gamma, args.warmup_factor, args.warmup_iters, args.warmup_methods) trainer = Trainer(model, criterion_z, trainvallabel) flag = 1 best_top1 = -1 # Start training for epoch in range(start_epoch, args.epochs): scheduler.step() id_loss = trainer.train(epoch, train_loader, optimizer_ft) save_checkpoint( { 'model': model.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, False, epoch, args.logs_dir, fpath='checkpoint.pth.tar') if epoch < 100: continue if not epoch % 10 == 0: continue top1 = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, metric) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'model': model.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, epoch, args.logs_dir, fpath='checkpoint.pth.tar') print('Test with best model:') print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'.format( epoch, top1, best_top1, ' *' if is_best else '')) checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) model.load_state_dict(checkpoint['model']) metric.train(model, train_loader) evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery, metric) print(args)
model.eval() modelTest.eval() if torch.cuda.is_available(): model = model.cuda() modelTest = modelTest.cuda() features, _ = extract_features(model, mix_loader, print_freq=10) features = torch.stack([features[f] for f, _, _ in mix_loader.dataset.dataset]) ncentroids = 512 fDim = features.shape[1] cluster = faiss.Kmeans(fDim, ncentroids, niter=20, gpu=True) cluster.train(features.cpu().numpy()) centroids = torch.from_numpy(cluster.centroids).cuda().float() evaluator = Evaluator(modelTest, args.print_freq) evaSrc = Evaluator(model, args.print_freq) # universal noise noise = torch.zeros((3, args.height, args.width)).cuda() normalize = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) noise.requires_grad = True MAX_EPS = args.max_eps / 255.0 optimizer = MI_SGD( [ {"params": [noise], "lr": MAX_EPS / 10, "momentum": 1, "sign": True} ], max_eps=MAX_EPS, )
def main(args): random.seed(args.seed) np.random.seed(1) torch.manual_seed(1) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # params params = {} if args.model_type == 'masks': params['num_m_features'] = args.num_m_features params['masks'] = args.masks else: print('unkrown model type.') return # Create data loaders if args.height is None or args.width is None: args.height, args.width = (256, 256) dataset, num_classes, random_train_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.logs_dir, args.model_type, params, args.height, args.width, args.crop_height, args.crop_width, args.batch_size, args.workers, args.combine_trainval) # Create model model = models.create(args.arch, num_classes=num_classes, params=params) # Load from checkpoint start_epoch = best_top1 = best_mAP = 0 if args.weights: checkpoint = load_checkpoint(args.weights) model_dict = model.state_dict() checkpoint_load = { k: v for k, v in (checkpoint['state_dict']).items() if k in model_dict } model_dict.update(checkpoint_load) model.load_state_dict(model_dict) if args.resume and not args.weights: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] print("=> Start epoch {} best top1 {:.1%}".format( start_epoch, best_top1)) model = nn.DataParallel(model, [0, 1, 2, 3]).cuda() # Criterion criterion = TripletLoss().cuda() # Optimizer base_params = [] for name, p in model.module.base.named_parameters(): base_params.append(p) base_param_ids = set(map(id, model.module.base.parameters())) new_params = [p for p in model.parameters() if id(p) not in base_param_ids] if args.model_type == 'masks': param_groups = [{ 'params': base_params, 'lr_mult': args.lrm }, { 'params': new_params, 'lr_mult': 1.0 }] else: print('unkrown model type.') return optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Trainer trainer = Trainer(model, criterion, num_classes, args.logs_dir) # Evaluator evaluator = Evaluator(model) # Schedule learning rate def adjust_lr(epoch): step_size = args.step_size if epoch < step_size: lr = args.lr elif epoch >= step_size and epoch < args.epochs: lr = args.lr * 0.1 for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) return lr # Start training for epoch in range(start_epoch, args.epochs): lr = adjust_lr(epoch) if epoch < args.warm_up_ep: trainer.train(epoch, random_train_loader, optimizer, lr, True, args.warm_up_ep) else: trainer.train(epoch, random_train_loader, optimizer, lr, False, args.warm_up_ep) if epoch < args.start_save: continue if epoch % 10 == 9: print('Epoch: [%d]' % epoch) top1, mAP = evaluator.evaluate(test_loader, dataset.query, dataset.gallery) is_best = mAP > best_mAP best_mAP = max(mAP, best_mAP) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_mAP': mAP, }, is_best, fpath=osp.join(args.logs_dir, 'model_best.pth.tar')) if epoch == args.epochs - 1: save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_mAP': mAP, }, True, fpath=osp.join(args.logs_dir, 'last.pth.tar'))
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) dataset, num_classes, train_loader, query_loader, gallery_loader = \ get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers, ) # Create model model = models.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=num_classes, cut_at_pooling=False, FCN=True) # Load from checkpoint start_epoch = best_mAP = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model_dict = model.state_dict() checkpoint_load = { k: v for k, v in (checkpoint['state_dict']).items() if k in model_dict } model_dict.update(checkpoint_load) model.load_state_dict(model_dict) # model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] best_top1 = checkpoint['best_mAP'] print("=> Start epoch {} best mAP {:.1%}".format( start_epoch, best_top1)) model = nn.DataParallel(model).cuda() # Evaluator evaluator = Evaluator(model) evaluator_6stripes = Evaluator_6stripes(model) if args.evaluate: print("Test:") #evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) evaluator_6stripes.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) return # Criterion criterion = nn.CrossEntropyLoss().cuda() # Optimizer if hasattr(model.module, 'base'): base_param_ids = set(map(id, model.module.base.parameters())) new_params = [ p for p in model.parameters() if id(p) not in base_param_ids ] param_groups = [{ 'params': model.module.base.parameters(), 'lr_mult': 0.1 }, { 'params': new_params, 'lr_mult': 1.0 }] else: param_groups = model.parameters() optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Trainer trainer = Trainer_6stripes(model, criterion, 0, 0, SMLoss_mode=0) # Schedule learning rate def adjust_lr(epoch): step_size = 60 if args.arch == 'inception' else args.step_size lr = args.lr * (0.1**(epoch // step_size)) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, train_loader, optimizer) def save_model(model_ema, is_best, best_mAP, mid): save_checkpoint( { 'state_dict': model_ema.state_dict(), 'epoch': epoch + 1, 'best_mAP': best_mAP, }, is_best, fpath=osp.join(args.logs_dir, 'model' + str(mid) + '_checkpoint.pth.tar')) if ((epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1)): R1, mAP = evaluator.evaluate(query_loader, gallery_loader, dataset.query, dataset.gallery) is_best = mAP > best_mAP if is_best: best_mAP = mAP save_model(model, is_best, best_mAP, epoch) print('\n * Finished epoch {:3d} mAP: {:5.1%} best: {:5.1%}{}\n'. format(epoch, mAP, best_mAP, ' *' if is_best else ''))
def main(args): print(args) np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders if args.big_height is None or args.big_width is None or args.target_height is None or args.target_width is None: args.big_height, args.big_width, args.target_height, args.target_width = ( 256, 256, 224, 224) dataset, num_classes, train_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.big_height, args.big_width, args.target_height, args.target_width, args.batch_size, args.num_instances, args.workers, args.combine_trainval) # Create model model = models.create(args.arch, num_classes=num_classes, num_features=args.features, is_cls=args.is_cls) # Load from checkpoint start_epoch = best = 0 if args.weights: #model_dict = model.state_dict() #checkpoint_load = {k: v for k, v in (checkpoint['state_dict']).items() if k in model_dict} #model_dict.update(checkpoint_load) #model.load_state_dict(model_dict) if args.arch == 'cross_trihard_senet101': model.base.load_param(args.weights) else: checkpoint = load_checkpoint(args.weights) del (checkpoint['fc.weight']) del (checkpoint['fc.bias']) model.base.load_state_dict(checkpoint) if args.resume and not args.weights: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] print("=> Start epoch {} best top1 {:.1%}".format(start_epoch, best)) model = nn.DataParallel(model).cuda() # Evaluator evaluator = Evaluator(model) if args.evaluate: print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery) return # Criterion ranking_loss = nn.MarginRankingLoss(margin=args.margin).cuda() if args.multi_attribute == 1: criterion = { 'MultiAttributeLoss': MultiAttributeLoss(is_cls=args.is_cls).cuda(), \ 'trihard': TripletLoss(ranking_loss).cuda() } else: criterion = { 'TypeAttributeLoss': TypeAttributeLoss(is_cls=args.is_cls).cuda(), \ 'trihard': TripletLoss(ranking_loss).cuda() } # Optimizer if hasattr(model.module, 'base'): base_params = [] base_bn_params = [] for name, p in model.module.base.named_parameters(): if 'bn' in name: base_bn_params.append(p) else: base_params.append(p) base_param_ids = set(map(id, model.module.base.parameters())) new_params = [ p for p in model.parameters() if id(p) not in base_param_ids ] param_groups = [{ 'params': base_params, 'lr_mult': 0.1 }, { 'params': base_bn_params, 'lr_mult': 0.1 }, { 'params': new_params, 'lr_mult': args.lr_mult }] else: param_groups = model.parameters() if args.optimizer == 0: print('SGD') optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) else: print('Adam') optimizer = torch.optim.Adam(params=param_groups, lr=args.lr, weight_decay=args.weight_decay) # Trainer if args.multi_attribute == 1: print('Multi Attribute') trainer = Multi_Attribute_Trainer( model, criterion, metric_loss_weight=args.metric_loss_weight, sub_task_loss_weight=args.sub_task_loss_weight) else: print('Type Attribute') trainer = Type_Attribute_Trainer( model, criterion, metric_loss_weight=args.metric_loss_weight, sub_task_loss_weight=args.sub_task_loss_weight) # Schedule learning rate def adjust_lr(epoch): step_size, step_size2, step_size3 = args.step_size, args.step_size2, args.step_size3 if epoch <= step_size: lr = args.lr elif epoch <= step_size2: lr = args.lr * 0.1 elif epoch <= step_size3: lr = args.lr * 0.01 else: lr = args.lr * 0.001 for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) return lr # Start training for epoch in range(start_epoch + 1, args.epochs + 1): lr = adjust_lr(epoch) trainer.train(epoch, train_loader, optimizer, lr, warm_up=True, warm_up_ep=args.warm_up_ep) if epoch % args.epoch_inter == 0 or epoch >= args.dense_evaluate: tmp_res = evaluator.evaluate(test_loader, dataset.query, dataset.gallery, is_attribute=True) if tmp_res > best and epoch >= args.start_save: best = tmp_res save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch, }, False, fpath=osp.join(args.logs_dir, 'pass%d.pth.tar' % (epoch)))
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) dataset, num_classes, train_loader, val_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.height, args.width, args.batch_size, args.workers, args.combine_trainval) # Create model model = models.create(args.arch, num_features=args.features, norm=True, dropout=args.dropout) # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] best_top1 = checkpoint['best_top1'] print("=> Start epoch {} best top1 {:.1%}".format( start_epoch, best_top1)) model = nn.DataParallel(model).cuda() # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model) if args.evaluate: metric.train(model, train_loader) print("Validation:") evaluator.evaluate(val_loader, dataset.val, dataset.val, metric) print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric) return # Criterion criterion = OIMLoss(model.module.num_features, num_classes, scalar=args.oim_scalar, momentum=args.oim_momentum).cuda() # Optimizer if hasattr(model.module, 'base'): base_param_ids = set(map(id, model.module.base.parameters())) new_params = [ p for p in model.parameters() if id(p) not in base_param_ids ] param_groups = [{ 'params': model.module.base.parameters(), 'lr_mult': 0.1 }, { 'params': new_params, 'lr_mult': 1.0 }] else: param_groups = model.parameters() optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Trainer trainer = Trainer(model, criterion) # Schedule learning rate def adjust_lr(epoch): step_size = 60 if args.arch == 'inception' else 40 lr = args.lr * (0.1**(epoch // step_size)) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, train_loader, optimizer) top1 = evaluator.evaluate(val_loader, dataset.val, dataset.val) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(epoch, top1, best_top1, ' *' if is_best else '')) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) model.module.load_state_dict(checkpoint['state_dict']) metric.train(model, train_loader) evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Create data loaders assert args.num_instances > 1, "num_instances should be greater than 1" assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) # get source data src_dataset, src_extfeat_loader = \ get_source_data(args.src_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # get target data tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \ get_data(args.tgt_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # Create model # Hacking here to let the classifier be the number of source ids if args.src_dataset == 'dukemtmc': model = models.create(args.arch, num_classes=632, pretrained=False) elif args.src_dataset == 'market1501': model = models.create(args.arch, num_classes=676, pretrained=False) else: raise RuntimeError( 'Please specify the number of classes (ids) of the network.') # Load from checkpoint if args.resume: print( 'Resuming checkpoints from finetuned model on another dataset...\n' ) checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict'], strict=False) else: raise RuntimeWarning('Not using a pre-trained model.') model = nn.DataParallel(model).cuda() # Distance metric # metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model, print_freq=args.print_freq) print( "Test with the original model trained on source domain (direct transfer):" ) rank_score_best = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) best_map = rank_score_best.map #market1501[0]-->rank-1 if args.evaluate: return # Criterion criterion = [ TripletLoss(args.margin, args.num_instances).cuda(), TripletLoss(args.margin, args.num_instances).cuda(), AccumulatedLoss(args.margin, args.num_instances).cuda(), nn.CrossEntropyLoss().cuda() ] # Optimizer optimizer = torch.optim.SGD( model.parameters(), lr=args.lr, momentum=0.9, ) # training stage transformer on input images normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transformer = T.Compose([ T.Resize((args.height, args.width)), T.RandomHorizontalFlip(), T.ToTensor(), normalizer, T.RandomErasing(probability=0.5, sh=0.2, r1=0.3) ]) # Start training for iter_n in range(args.iteration): if args.lambda_value == 0: source_features = 0 #this value controls the usage of source data else: # get source datas' feature source_features, _ = extract_features(model, src_extfeat_loader, print_freq=args.print_freq) # synchronization feature order with src_dataset.train source_features = torch.cat([ source_features[f].unsqueeze(0) for f, _, _ in src_dataset.train ], 0) # extract training images' features print('Iteration {}: Extracting Target Dataset Features...'.format( iter_n + 1)) target_features, _ = extract_features(model, tgt_extfeat_loader, print_freq=args.print_freq) # synchronization feature order with dataset.train target_features = torch.cat([ target_features[f].unsqueeze(0) for f, _, _ in tgt_dataset.trainval ], 0) # calculate distance and rerank result print('Calculating feature distances...') target_features = target_features.numpy() rerank_dist = re_ranking(source_features, target_features, lambda_value=args.lambda_value) if iter_n == 0: # DBSCAN cluster tri_mat = np.triu(rerank_dist, 1) # tri_mat.dim=2 tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1 tri_mat = np.sort(tri_mat, axis=None) top_num = np.round(args.rho * tri_mat.size).astype(int) eps = tri_mat[:top_num].mean() print('eps in cluster: {:.3f}'.format(eps)) cluster = DBSCAN(eps=eps, min_samples=4, metric='precomputed', n_jobs=8) # HDBSCAN cluster import hdbscan cluster_hdbscan = hdbscan.HDBSCAN(min_cluster_size=10, min_samples=4, metric='precomputed') # select & cluster images as training set of this epochs print('Clustering and labeling...') if args.use_hdbscan_clustering: print( 'Use the better chlustering algorithm HDBSCAN for clustering' ) labels = cluster_hdbscan.fit_predict(rerank_dist) else: print('Use DBSCAN for clustering') labels = cluster.fit_predict(rerank_dist) num_ids = len(set(labels)) - 1 print('Only do once, Iteration {} have {} training ids'.format( iter_n + 1, num_ids)) # generate new dataset new_dataset = [] for (fname, _, _), label in zip(tgt_dataset.trainval, labels): if label == -1: continue # dont need to change codes in trainer.py _parsing_input function and sampler function after add 0 new_dataset.append((fname, label, 0)) print('Only do once, Iteration {} have {} training images'.format( iter_n + 1, len(new_dataset))) train_loader = DataLoader( Preprocessor_return_index(new_dataset, root=tgt_dataset.images_dir, transform=train_transformer), batch_size=args.batch_size, num_workers=4, sampler=RandomIdentitySampler(new_dataset, args.num_instances), pin_memory=True, drop_last=True) # init pseudo/fake labels, y_tilde in cvpr19's paper: new_label = np.zeros([len(new_dataset), num_ids]) # init y_tilde, let softmax(y_tilde) is noisy labels for index, (imgs, _, pids, _, index) in enumerate(train_loader): index = index.numpy() onehot = torch.zeros(pids.size(0), num_ids).scatter_(1, pids.view(-1, 1), 10.0) onehot = onehot.numpy() new_label[index, :] = onehot # Using clustered label to init the new classifier: classifier = nn.Linear(2048, num_ids, bias=False) classifier.apply(weights_init_classifier) classifier = nn.DataParallel(classifier).cuda() optimizer_cla = torch.optim.SGD(classifier.parameters(), lr=args.lr * 10, momentum=0.9) # train model with new generated dataset trainer = Trainer_with_learnable_label(model, classifier, criterion, print_freq=args.print_freq) evaluator = Evaluator(model, print_freq=args.print_freq) # Start training for epoch in range(args.epochs): trainer.train(epoch, train_loader, new_label, optimizer, optimizer_cla) # Evaluate rank_score = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) #Save the best ckpt: rank1 = rank_score.market1501[0] mAP = rank_score.map is_best_mAP = mAP > best_map best_map = max(mAP, best_map) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': iter_n + 1, 'best_mAP': best_map, # 'num_ids': num_ids, }, is_best_mAP, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print( '\n * Finished epoch {:3d} top1: {:5.1%} mAP: {:5.1%} best_mAP: {:5.1%}{}\n' .format(iter_n + 1, rank1, mAP, best_map, ' *' if is_best_mAP else '')) return (rank_score.map, rank_score.market1501[0])
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) dataset, num_classes, train_loader, val_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.height, args.width, args.batch_size, args.workers, args.combine_trainval) # Create model model = models.create(args.arch, num_features=args.features, norm=True, dropout=args.dropout) # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] best_top1 = checkpoint['best_top1'] print("=> Start epoch {} best top1 {:.1%}" .format(start_epoch, best_top1)) model = nn.DataParallel(model).cuda() # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model) if args.evaluate: metric.train(model, train_loader) print("Validation:") evaluator.evaluate(val_loader, dataset.val, dataset.val, metric) print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric) return # Criterion criterion = OIMLoss(model.module.num_features, num_classes, scalar=args.oim_scalar, momentum=args.oim_momentum).cuda() # Optimizer if hasattr(model.module, 'base'): base_param_ids = set(map(id, model.module.base.parameters())) new_params = [p for p in model.parameters() if id(p) not in base_param_ids] param_groups = [ {'params': model.module.base.parameters(), 'lr_mult': 0.1}, {'params': new_params, 'lr_mult': 1.0}] else: param_groups = model.parameters() optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Trainer trainer = Trainer(model, criterion) # Schedule learning rate def adjust_lr(epoch): step_size = 60 if args.arch == 'inception' else 40 lr = args.lr * (0.1 ** (epoch // step_size)) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, train_loader, optimizer) if epoch < args.start_save: continue top1 = evaluator.evaluate(val_loader, dataset.val, dataset.val) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint({ 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(epoch, top1, best_top1, ' *' if is_best else '')) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) model.module.load_state_dict(checkpoint['state_dict']) metric.train(model, train_loader) evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
def evaluate(self, query, gallery): test_loader = self.get_dataloader(list(set(query) | set(gallery)), training = False) evaluator = Evaluator(self.model) return evaluator.evaluate(test_loader, query, gallery)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True #######Tensorboard-logs########## logger = TensorLogger(args.Tensorlogs_dir) # Redirect print to both console and log file if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) # Create data loaders if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) dataset, _, _, val_loader, test_loader = \ get_data(args.dataset, args.split, args.data_dir, args.height, args.width, args.batch_size, args.workers, args.combine_trainval) trans_train_loader, num_classes = get_fake_data(args.trans_name, args.trans_data_txt, args.height, args.width, args.batch_size, args.workers) # Create model model = models.create(args.arch, dropout=0, num_classes=num_classes) # model = models.create(args.arch, num_features=1024, num_diff_features=args.features, # dropout=args.dropout, num_classes=num_classes, iden_pretrain=True) # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] best_top1 = checkpoint['best_top1'] print("=> Start epoch {} best top1 {:.1%}" .format(start_epoch, best_top1)) model = nn.DataParallel(model).cuda() # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model) if args.evaluate: metric.train(model, trans_train_loader) # print("Validation:") # evaluator.evaluate(val_loader, dataset.val, dataset.val, metric) print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric) return # Criterion criterion = nn.CrossEntropyLoss().cuda() # Optimizer if hasattr(model.module, 'base'): base_param_ids = set(map(id, model.module.base.parameters())) new_params = [p for p in model.parameters() if id(p) not in base_param_ids] param_groups = [ {'params': model.module.base.parameters(), 'lr_mult': 0.1}, {'params': new_params, 'lr_mult': 1.0}] else: param_groups = model.parameters() optimizer = torch.optim.SGD(param_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True) # Trainer trainer = Trainer(model, criterion) # Schedule learning rate def adjust_lr(epoch): # lr = args.lr if epoch <= 270 else \ # args.lr * (0.001 ** ((epoch - 270) / 135)) # lr = args.lr if epoch <= 100 else \ # args.lr * (0.001 ** ((epoch - 100) / 50.0)) lr = args.lr * 0.1**(epoch//args.lr_change_epochs) for g in optimizer.param_groups: g['lr'] = lr * g.get('lr_mult', 1) # Start training for epoch in range(start_epoch, args.epochs): adjust_lr(epoch) trainer.train(epoch, trans_train_loader, optimizer, logger) #######Tensorboard-logs########## # for tag, value in model.named_parameters(): # tag = tag.replace('.', '/') # try: # logger.histo_summary(tag, to_np(value), epoch * len(train_loader) + 1) # logger.histo_summary(tag + '/grad', to_np(value.grad), epoch * len(train_loader) + 1) # except AttributeError, e: # pass # for tag, value in criterion.named_parameters(): # tag = tag.replace('.', '/') # try: # logger.histo_summary(tag, to_np(value), epoch * len(train_loader) + 1) # logger.histo_summary(tag + '/grad', to_np(value.grad), epoch * len(train_loader) + 1) # except AttributeError, e: # pass ################################# top1 = evaluator.evaluate(val_loader, dataset.val, dataset.val) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint({ 'state_dict': model.module.state_dict(), 'epoch': epoch + 1, 'best_top1': best_top1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(epoch, top1, best_top1, ' *' if is_best else '')) # Final test print('Test with best model:') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) model.module.load_state_dict(checkpoint['state_dict']) metric.train(model, trans_train_loader) evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric)
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Create data loaders assert args.num_instances > 1, "num_instances should be greater than 1" assert args.batch_size % args.num_instances == 0, \ 'num_instances should divide batch_size' if args.height is None or args.width is None: args.height, args.width = (144, 56) if args.arch == 'inception' else \ (256, 128) # get_source_data src_dataset, src_extfeat_loader = \ get_source_data(args.src_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # get_target_data tgt_dataset, num_classes, tgt_extfeat_loader, test_loader = \ get_data(args.tgt_dataset, args.data_dir, args.height, args.width, args.batch_size, args.workers) # Create model # Hacking here to let the classifier be the number of source ids if args.src_dataset == 'dukemtmc': model = models.create(args.arch, num_classes=632, pretrained=False) elif args.src_dataset == 'market1501': model = models.create(args.arch, num_classes=676, pretrained=False) else: raise RuntimeError( 'Please specify the number of classes (ids) of the network.') # Load from checkpoint start_epoch = best_top1 = 0 if args.resume: print( 'Resuming checkpoints from finetuned model on another dataset...\n' ) checkpoint = load_checkpoint(args.resume) model.load_state_dict(checkpoint['state_dict'], strict=False) else: raise RuntimeWarning('Not using a pre-trained model.') model = nn.DataParallel(model).cuda() # Distance metric metric = DistanceMetric(algorithm=args.dist_metric) # Evaluator evaluator = Evaluator(model, print_freq=args.print_freq) print("Test with the original model trained on source domain:") evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) if args.evaluate: return # Criterion criterion = [] criterion.append( TripletLoss(margin=args.margin, num_instances=args.num_instances).cuda()) criterion.append( TripletLoss(margin=args.margin, num_instances=args.num_instances).cuda()) # Optimizer optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9) # training stage transformer on input images normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transformer = T.Compose([ T.Resize((args.height, args.width)), T.RandomHorizontalFlip(), T.ToTensor(), normalizer, T.RandomErasing(probability=0.5, sh=0.2, r1=0.3) ]) # Start training for iter_n in range(args.iteration): if args.lambda_value == 0: source_features = 0 else: # get source datas' feature source_features, _ = extract_features(model, src_extfeat_loader, print_freq=args.print_freq) # synchronization feature order with src_dataset.train source_features = torch.cat([ source_features[f].unsqueeze(0) for f, _, _ in src_dataset.train ], 0) # extract training images' features print('Iteration {}: Extracting Target Dataset Features...'.format( iter_n + 1)) target_features, _ = extract_features(model, tgt_extfeat_loader, print_freq=args.print_freq) # synchronization feature order with dataset.train target_features = torch.cat( [target_features[f].unsqueeze(0) for f, _, _ in tgt_dataset.train], 0) # calculate distance and rerank result print('Calculating feature distances...') target_features = target_features.numpy() rerank_dist = re_ranking(source_features, target_features, lambda_value=args.lambda_value) if iter_n == 0: # DBSCAN cluster tri_mat = np.triu(rerank_dist, 1) # tri_mat.dim=2 tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1 tri_mat = np.sort(tri_mat, axis=None) top_num = np.round(args.rho * tri_mat.size).astype(int) eps = tri_mat[:top_num].mean() print('eps in cluster: {:.3f}'.format(eps)) cluster = DBSCAN(eps=eps, min_samples=4, metric='precomputed', n_jobs=8) # select & cluster images as training set of this epochs print('Clustering and labeling...') labels = cluster.fit_predict(rerank_dist) num_ids = len(set(labels)) - 1 print('Iteration {} have {} training ids'.format(iter_n + 1, num_ids)) # generate new dataset new_dataset = [] for (fname, _, _), label in zip(tgt_dataset.trainval, labels): if label == -1: continue # dont need to change codes in trainer.py _parsing_input function and sampler function after add 0 new_dataset.append((fname, label, 0)) print('Iteration {} have {} training images'.format( iter_n + 1, len(new_dataset))) train_loader = DataLoader(Preprocessor(new_dataset, root=tgt_dataset.images_dir, transform=train_transformer), batch_size=args.batch_size, num_workers=4, sampler=RandomIdentitySampler( new_dataset, args.num_instances), pin_memory=True, drop_last=True) # train model with new generated dataset trainer = Trainer(model, criterion, print_freq=args.print_freq) evaluator = Evaluator(model, print_freq=args.print_freq) # Start training for epoch in range(args.epochs): trainer.train(epoch, train_loader, optimizer) # Evaluate rank1 = evaluator.evaluate(test_loader, tgt_dataset.query, tgt_dataset.gallery) save_checkpoint( { 'state_dict': model.module.state_dict(), 'epoch': iter_n + 1, 'num_ids': num_ids, }, True, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print('\n * Finished epoch {:3d} rank1: {:5.1%} \n'.format( iter_n + 1, rank1))