def create_model(args): model_1 = models_pre_a_s2.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=args.num_clusters) model_2 = models_pre_a_s2.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=args.num_clusters) model_1_ema = models_pre_a_s2.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=args.num_clusters) model_2_ema = models_pre_a_s2.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=args.num_clusters) model_1.cuda() model_2.cuda() model_1_ema.cuda() model_2_ema.cuda() model_1 = nn.DataParallel(model_1) model_2 = nn.DataParallel(model_2) model_1_ema = nn.DataParallel(model_1_ema) model_2_ema = nn.DataParallel(model_2_ema) initial_weights = load_checkpoint(args.init_1) copy_state_dict(initial_weights['state_dict'], model_1) copy_state_dict(initial_weights['state_dict'], model_1_ema) model_1_ema.module.classifier.weight.data.copy_(model_1.module.classifier.weight.data) initial_weights = load_checkpoint(args.init_2) copy_state_dict(initial_weights['state_dict'], model_2) copy_state_dict(initial_weights['state_dict'], model_2_ema) model_2_ema.module.classifier.weight.data.copy_(model_2.module.classifier.weight.data) for param in model_1_ema.parameters(): param.detach_() for param in model_2_ema.parameters(): param.detach_() return model_1, model_2, model_1_ema, model_2_ema
def main_worker(args): cudnn.benchmark = True log_dir = osp.dirname(args.resume) sys.stdout = Logger(osp.join(log_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) # Create data loaders dataset_target, test_loader_target = \ get_data(args.dataset_target, args.data_dir, args.height, args.width, args.batch_size, args.workers) # Create model model = models_post_a_s1.create(args.arch, pretrained=False, num_features=args.features, dropout=args.dropout, num_classes=0) model.cuda() model = nn.DataParallel(model) # Load from checkpoint checkpoint = load_checkpoint(args.resume) copy_state_dict(checkpoint['state_dict'], model) start_epoch = checkpoint['epoch'] best_mAP = checkpoint['best_mAP'] print("=> Checkpoint of epoch {} best mAP {:.1%}".format( start_epoch, best_mAP)) # Evaluator evaluator = Evaluator(model) print("Test on the target domain of {}:".format(args.dataset_target)) evaluator.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery, cmc_flag=True, rerank=args.rerank) return
def main_worker(args): global best_mAP cudnn.benchmark = True sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) print("==========\nArgs:{}\n==========".format(args)) # Create data loaders iters = args.iters if (args.iters > 0) else None dataset_target = get_data(args.dataset_target, args.data_dir) test_loader_target = get_test_loader(dataset_target, args.height, args.width, args.batch_size, args.workers) cluster_loader = get_test_loader(dataset_target, args.height, args.width, args.batch_size, args.workers, testset=dataset_target.train) # Create model model_1, model_2, model_1_ema, model_2_ema = create_model(args) # Evaluator evaluator_1_ema = Evaluator(model_1_ema) evaluator_2_ema = Evaluator(model_2_ema) for epoch in range(args.epochs): dict_f, _ = extract_features(model_1_ema, cluster_loader, print_freq=50) cf_1 = torch.stack(list(dict_f.values())).numpy() dict_f, _ = extract_features(model_2_ema, cluster_loader, print_freq=50) cf_2 = torch.stack(list(dict_f.values())).numpy() cf = (cf_1 + cf_2) / 2 print('\n Clustering into {} classes \n'.format(args.num_clusters)) km = KMeans(n_clusters=args.num_clusters, random_state=args.seed, n_jobs=2).fit(cf) model_1.module.classifier.weight.data.copy_( torch.from_numpy(normalize(km.cluster_centers_, axis=1)).float().cuda()) model_2.module.classifier.weight.data.copy_( torch.from_numpy(normalize(km.cluster_centers_, axis=1)).float().cuda()) model_1_ema.module.classifier.weight.data.copy_( torch.from_numpy(normalize(km.cluster_centers_, axis=1)).float().cuda()) model_2_ema.module.classifier.weight.data.copy_( torch.from_numpy(normalize(km.cluster_centers_, axis=1)).float().cuda()) target_label = km.labels_ # change pseudo labels for i in range(len(dataset_target.train)): dataset_target.train[i] = list(dataset_target.train[i]) dataset_target.train[i][1] = int(target_label[i]) dataset_target.train[i] = tuple(dataset_target.train[i]) train_loader_target = get_train_loader(dataset_target, args.height, args.width, args.batch_size, args.workers, args.num_instances, iters) # Optimizer params = [] for key, value in model_1.named_parameters(): if not value.requires_grad: continue if 'non_local' in key: params += [{ "params": [value], "lr": args.lr, "weight_decay": args.weight_decay }] for key, value in model_2.named_parameters(): if not value.requires_grad: continue if 'non_local' in key: params += [{ "params": [value], "lr": args.lr, "weight_decay": args.weight_decay }] optimizer = torch.optim.Adam(params) # Trainer trainer = MMTTrainer(model_1, model_2, model_1_ema, model_2_ema, num_cluster=args.num_clusters, alpha=args.alpha) train_loader_target.new_epoch() trainer.train(epoch, train_loader_target, optimizer, ce_soft_weight=args.soft_ce_weight, tri_soft_weight=args.soft_tri_weight, print_freq=args.print_freq, train_iters=len(train_loader_target)) 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)): mAP_1 = evaluator_1_ema.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery, cmc_flag=False) mAP_2 = evaluator_2_ema.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery, cmc_flag=False) is_best = (mAP_1 > best_mAP) or (mAP_2 > best_mAP) best_mAP = max(mAP_1, mAP_2, best_mAP) save_model(model_1_ema, (is_best and (mAP_1 > mAP_2)), best_mAP, 1) save_model(model_2_ema, (is_best and (mAP_1 <= mAP_2)), best_mAP, 2) print( '\n * Finished epoch {:3d} model no.1 mAP: {:5.1%} model no.2 mAP: {:5.1%} best: {:5.1%}{}\n' .format(epoch, mAP_1, mAP_2, best_mAP, ' *' if is_best else '')) print('Test on the best model.') checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar')) model_1_ema.load_state_dict(checkpoint['state_dict']) evaluator_1_ema.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery, cmc_flag=True)
def main_worker(args): global start_epoch, best_mAP cudnn.benchmark = True if not args.evaluate: sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) else: log_dir = osp.dirname(args.resume) sys.stdout = Logger(osp.join(log_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) # Create data loaders iters = args.iters if (args.iters > 0) else None dataset_source, num_classes, train_loader_source, test_loader_source = \ get_data(args.dataset_source, args.data_dir, args.height, args.width, args.batch_size, args.workers, args.num_instances, iters) dataset_target, _, train_loader_target, test_loader_target = \ get_data(args.dataset_target, args.data_dir, args.height, args.width, args.batch_size, args.workers, 0, iters) # Create model model = models_source.create(args.arch, num_features=args.features, dropout=args.dropout, num_classes=num_classes) model.cuda() model = nn.DataParallel(model) # Load from checkpoint if args.resume: checkpoint = load_checkpoint(args.resume) copy_state_dict(checkpoint['state_dict'], model) start_epoch = checkpoint['epoch'] best_mAP = checkpoint['best_mAP'] print("=> Start epoch {} best mAP {:.1%}".format( start_epoch, best_mAP)) # Evaluator evaluator = Evaluator(model) if args.evaluate: print("Test on source domain:") evaluator.evaluate(test_loader_source, dataset_source.query, dataset_source.gallery, cmc_flag=True, rerank=args.rerank) print("Test on target domain:") evaluator.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery, cmc_flag=True, rerank=args.rerank) return params = [] for key, value in model.named_parameters(): if not value.requires_grad: continue params += [{ "params": [value], "lr": args.lr, "weight_decay": args.weight_decay }] optimizer = torch.optim.Adam(params) lr_scheduler = WarmupMultiStepLR(optimizer, args.milestones, gamma=0.1, warmup_factor=0.01, warmup_iters=args.warmup_step) # Trainer trainer = PreTrainer(model, num_classes, margin=args.margin) # Start training for epoch in range(start_epoch, args.epochs): lr_scheduler.step() train_loader_source.new_epoch() train_loader_target.new_epoch() trainer.train(epoch, train_loader_source, train_loader_target, optimizer, train_iters=len(train_loader_source), print_freq=args.print_freq) if ((epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1)): _, mAP = evaluator.evaluate(test_loader_source, dataset_source.query, dataset_source.gallery, cmc_flag=True) is_best = mAP > best_mAP best_mAP = max(mAP, best_mAP) save_checkpoint( { 'state_dict': model.state_dict(), 'epoch': epoch + 1, 'best_mAP': best_mAP, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) print( '\n * Finished epoch {:3d} source mAP: {:5.1%} best: {:5.1%}{}\n' .format(epoch, mAP, best_mAP, ' *' if is_best else '')) print("Test on target domain:") evaluator.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery, cmc_flag=True, rerank=args.rerank)