def get_train_loader(dataset, height, width, batch_size, workers, num_instances, iters, trainset=None): normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transformer = T.Compose([ T.Resize((height, width), interpolation=3), T.RandomHorizontalFlip(p=0.5), T.Pad(10), T.RandomCrop((height, width)), T.ToTensor(), normalizer, T.RandomErasing(probability=0.5, mean=[0.485, 0.456, 0.406]) ]) train_set = dataset.train if trainset is None else trainset rmgs_flag = num_instances > 0 if rmgs_flag: sampler = RandomMultipleGallerySampler(train_set, num_instances) else: sampler = None train_loader = IterLoader( DataLoader(Preprocessor(train_set, root=dataset.images_dir, transform=train_transformer, mutual=False), batch_size=batch_size, num_workers=workers, sampler=sampler, shuffle=not rmgs_flag, pin_memory=True, drop_last=True), length=iters) return train_loader
def get_data(name, data_dir, height, width, batch_size, workers, num_instances, iters=200): root = osp.join(data_dir, name) dataset = datasets.create(name, root) normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_set = dataset.train num_classes = dataset.num_train_pids train_transformer = T.Compose([ T.Resize((height, width), interpolation=3), T.RandomHorizontalFlip(p=0.5), T.Pad(10), T.RandomCrop((height, width)), T.ToTensor(), normalizer ]) test_transformer = T.Compose([ T.Resize((height, width), interpolation=3), T.ToTensor(), normalizer ]) rmgs_flag = num_instances > 0 if rmgs_flag: sampler = RandomMultipleGallerySampler(train_set, num_instances) else: sampler = None train_loader = IterLoader( DataLoader(Preprocessor(train_set, root=dataset.images_dir, transform=train_transformer), batch_size=batch_size, num_workers=workers, sampler=sampler, shuffle=not rmgs_flag, pin_memory=True, drop_last=True), length=iters) test_loader = DataLoader( Preprocessor(list(set(dataset.query) | set(dataset.gallery)), root=dataset.images_dir, transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) return dataset, num_classes, train_loader, test_loader
def get_data(data_dir, source, target, height, width, batch_size, num_instance=2, workers=8): dataset = DA(data_dir, source, target) normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) num_classes = dataset.num_train_ids train_transformer = T.Compose([ T.Resize((256, 128), interpolation=3), T.Pad(10), T.RandomCrop((256,128)), T.RandomHorizontalFlip(0.5), T.RandomRotation(5), T.ToTensor(), normalizer, ]) test_transformer = T.Compose([ T.Resize((256, 128), interpolation=3), T.ToTensor(), normalizer, ]) source_train_loader = DataLoader( Preprocessor_occluded(dataset.source_train, root=osp.join(dataset.source_images_dir, dataset.source_train_path), transform=train_transformer, train=True), batch_size=batch_size, num_workers=workers, sampler=IdentitySampler(dataset.source_train, num_instance), pin_memory=True, drop_last=True) query_loader = DataLoader( Preprocessor_occluded(dataset.query, root=osp.join(dataset.target_images_dir, dataset.query_path), transform=test_transformer), batch_size=42, num_workers=workers, shuffle=False, pin_memory=True) gallery_loader = DataLoader( Preprocessor_occluded(dataset.gallery, root=osp.join(dataset.target_images_dir, dataset.gallery_path), transform=test_transformer), batch_size=42, num_workers=workers, shuffle=False, pin_memory=True) return dataset, num_classes, source_train_loader, query_loader, gallery_loader
def main(args): np.random.seed(args.seed) torch.manual_seed(args.seed) cudnn.benchmark = True # Redirect print to both console and log file sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) print(args) shutil.copy(sys.argv[0], osp.join(args.logs_dir, osp.basename(sys.argv[0]))) # Create data loaders 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 * 8, args.workers, ) # Create model model = models.create("ft_net_inter", num_classes=num_classes, stride=args.stride) # Load from checkpoint start_epoch = 0 best_top1 = 0 top1 = 0 is_best = False if args.checkpoint is not None: if args.evaluate: checkpoint = load_checkpoint(args.checkpoint) param_dict = model.state_dict() for k, v in checkpoint['state_dict'].items(): if 'model' in k: param_dict[k] = v model.load_state_dict(param_dict) else: model.model.load_param(args.checkpoint) model = model.cuda() # Distance metric metric = None # Evaluator evaluator = Evaluator(model, use_cpu=args.use_cpu) if args.evaluate: print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric) return train_transformer = [ T.Resize((args.height, args.width), interpolation=3), T.RandomHorizontalFlip(), T.Pad(10), T.RandomCrop((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), T.RandomErasing(probability=0.5), ] train_transformer = T.Compose(train_transformer) for cluster_epoch in range(args.cluster_epochs): # -------------------------Stage 1 intra camera training-------------------------- # Cluster and generate new dataset and model cluster_result = get_intra_cam_cluster_result(model, train_loader, args.class_number_stage1, args.linkage) cluster_datasets = [ datasets.create("cluster", osp.join(args.data_dir, args.dataset), cluster_result[cam_id], cam_id) for cam_id in cluster_result.keys() ] cluster_dataloaders = [ DataLoader(Preprocessor(dataset.train_set, root=dataset.images_dir, transform=train_transformer), batch_size=args.batch_size, num_workers=args.workers, shuffle=True, pin_memory=False, drop_last=True) for dataset in cluster_datasets ] param_dict = model.model.state_dict() model = models.create("ft_net_intra", num_classes=[ args.class_number_stage1 for cam_id in cluster_result.keys() ], stride=args.stride) model_param_dict = model.model.state_dict() for k, v in model_param_dict.items(): if k in param_dict.keys(): model_param_dict[k] = param_dict[k] model.model.load_state_dict(model_param_dict) model = model.cuda() criterion = nn.CrossEntropyLoss().cuda() # Optimizer param_groups = make_params(model, args.lr, args.weight_decay) optimizer = torch.optim.SGD(param_groups, momentum=0.9) # Trainer trainer = IntraCameraTrainer(model, criterion, warm_up_epoch=args.warm_up) print("start training") # Start training for epoch in range(0, args.epochs_stage1): trainer.train( cluster_epoch, epoch, cluster_dataloaders, optimizer, print_freq=args.print_freq, ) # -------------------------------------------Stage 2 inter camera training----------------------------------- mix_rate = get_mix_rate(args.mix_rate, cluster_epoch, args.cluster_epochs, power=args.decay_factor) cluster_result = get_inter_cam_cluster_result(model, train_loader, args.class_number_stage2, args.linkage, mix_rate, use_cpu=args.use_cpu) cluster_dataset = datasets.create( "cluster", osp.join(args.data_dir, args.dataset), cluster_result, 0) cluster_dataloaders = DataLoader( Preprocessor(cluster_dataset.train_set, root=cluster_dataset.images_dir, transform=train_transformer), batch_size=args.batch_size_stage2, num_workers=args.workers, sampler=RandomIdentitySampler(cluster_dataset.train_set, args.batch_size_stage2, args.instances), pin_memory=False, drop_last=True) param_dict = model.model.state_dict() model = models.create("ft_net_inter", num_classes=args.class_number_stage2, stride=args.stride) model.model.load_state_dict(param_dict) model = model.cuda() # Criterion criterion_entropy = nn.CrossEntropyLoss().cuda() criterion_triple = TripletLoss(margin=args.margin).cuda() # Optimizer param_groups = make_params(model, args.lr * args.batch_size_stage2 / 32, args.weight_decay) optimizer = torch.optim.SGD(param_groups, momentum=0.9) # Trainer trainer = InterCameraTrainer( model, criterion_entropy, criterion_triple, warm_up_epoch=args.warm_up, ) print("start training") # Start training for epoch in range(0, args.epochs_stage2): trainer.train(cluster_epoch, epoch, cluster_dataloaders, optimizer, print_freq=args.print_freq) if (cluster_epoch + 1) % 5 == 0: evaluator = Evaluator(model, use_cpu=args.use_cpu) top1, mAP = evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric, return_mAP=True) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'state_dict': model.state_dict(), 'epoch': cluster_epoch + 1, 'best_top1': best_top1, 'cluster_epoch': cluster_epoch + 1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) if cluster_epoch == (args.cluster_epochs - 1): save_checkpoint( { 'state_dict': model.state_dict(), 'epoch': cluster_epoch + 1, 'best_top1': best_top1, 'cluster_epoch': cluster_epoch + 1, }, False, fpath=osp.join(args.logs_dir, 'latest.pth.tar')) print('\n * cluster_epoch: {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(cluster_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']) best_rank1, mAP = evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric, return_mAP=True)
def get_data(name, split_id, data_dir, height, width, batch_size, num_instances, workers, combine_trainval, flip_prob, padding, re_prob): root = osp.join(data_dir, name) print(root) dataset = datasets.create(name, root, split_id=split_id) normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) trainvallabel = dataset.trainvallabel train_set = dataset.trainval if combine_trainval else dataset.train num_classes = (dataset.num_trainval_ids if combine_trainval else dataset.num_train_ids) train_transformer = T.Compose([ T.Resize((height, width)), T.RandomHorizontalFlip(), T.Pad(padding), T.RandomCrop((height, width)), T.ToTensor(), normalizer, RandomErasing(probability=re_prob, mean=[0.485, 0.456, 0.406]) ]) # train_transformer = T.Compose([ # T.RandomSizedRectCrop(height, width), # T.RandomHorizontalFlip(), # T.ToTensor(), # normalizer, # ]) test_transformer = T.Compose([ T.Resize((height, width)), T.ToTensor(), normalizer, ]) val_loader = DataLoader(Preprocessor(dataset.val, root=dataset.images_dir, transform=test_transformer), batch_size=32, num_workers=workers, shuffle=False, pin_memory=True) query_loader = DataLoader(Preprocessor(list(set(dataset.query)), root=dataset.images_dir, transform=test_transformer), batch_size=32, num_workers=workers, sampler=CamSampler(list(set(dataset.query)), [2, 5]), shuffle=False, pin_memory=True) gallery_loader = DataLoader(Preprocessor(list(set(dataset.gallery)), root=dataset.images_dir, transform=test_transformer), batch_size=32, num_workers=workers, sampler=CamSampler(list(set(dataset.gallery)), [0, 1, 3, 4], 4), shuffle=False, pin_memory=True) train_loader = DataLoader(Preprocessor(train_set, root=dataset.images_dir, transform=train_transformer), batch_size=batch_size, num_workers=workers, sampler=RandomIdentitySampler( train_set, num_instances), pin_memory=True, drop_last=True) return dataset, num_classes, train_loader, trainvallabel, val_loader, query_loader, gallery_loader
def get_data(data_dir, height, width, batch_size, num_instances, re=0, workers=8): dataset = DA(data_dir) test_dataset = TotalData(data_dir) normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) num_classes = dataset.num_source_ids train_transformer = T.Compose([ T.Resize((256, 128), interpolation=3), T.Pad(10), T.RandomCrop((256,128)), T.RandomHorizontalFlip(0.5), T.RandomRotation(5), T.ColorJitter(brightness=(0.5, 2.0), saturation=(0.5, 2.0), hue=(-0.1, 0.1)), T.ToTensor(), normalizer, # T.RandomErasing(EPSILON=re), ]) test_transformer = T.Compose([ T.Resize((256, 128), interpolation=3), T.ToTensor(), normalizer, ]) # Train source_train_loader = DataLoader( Preprocessor(dataset.source_train, transform=train_transformer), batch_size=batch_size, num_workers=workers, # shuffle=True, pin_memory=True, drop_last=True) sampler=RandomIdentitySampler(dataset.source_train, batch_size, num_instances), pin_memory=True, drop_last=True) # Test grid_query_loader = DataLoader( Preprocessor(test_dataset.grid_query, root=osp.join(test_dataset.grid_images_dir, test_dataset.query_path), transform=test_transformer), batch_size=64, num_workers=4, shuffle=False, pin_memory=True) grid_gallery_loader = DataLoader( Preprocessor(test_dataset.grid_gallery, root=osp.join(test_dataset.grid_images_dir, test_dataset.gallery_path), transform=test_transformer), batch_size=64, num_workers=4, shuffle=False, pin_memory=True) prid_query_loader = DataLoader( Preprocessor(test_dataset.prid_query, root=osp.join(test_dataset.prid_images_dir, test_dataset.query_path), transform=test_transformer), batch_size=64, num_workers=4, shuffle=False, pin_memory=True) prid_gallery_loader = DataLoader( Preprocessor(test_dataset.prid_gallery, root=osp.join(test_dataset.prid_images_dir, test_dataset.gallery_path), transform=test_transformer), batch_size=64, num_workers=4, shuffle=False, pin_memory=True) viper_query_loader = DataLoader( Preprocessor(test_dataset.viper_query, root=osp.join(test_dataset.viper_images_dir, test_dataset.query_path), transform=test_transformer), batch_size=64, num_workers=4, shuffle=False, pin_memory=True) viper_gallery_loader = DataLoader( Preprocessor(test_dataset.viper_gallery, root=osp.join(test_dataset.viper_images_dir, test_dataset.gallery_path), transform=test_transformer), batch_size=64, num_workers=4, shuffle=False, pin_memory=True) ilid_query_loader = DataLoader( Preprocessor(test_dataset.ilid_query, root=osp.join(test_dataset.ilid_images_dir, "images"), transform=test_transformer), batch_size=64, num_workers=4, shuffle=False, pin_memory=True) ilid_gallery_loader = DataLoader( Preprocessor(test_dataset.ilid_gallery, root=osp.join(test_dataset.ilid_images_dir, "images"), transform=test_transformer), batch_size=64, num_workers=4, shuffle=False, pin_memory=True) return 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
def get_data(dataname, data_dir, model, matcher, save_path, args): root = osp.join(data_dir, dataname) dataset = datasets.create(dataname, root, combine_all=args.combine_all) num_classes = dataset.num_train_ids train_transformer = T.Compose([ T.Resize((args.height, args.width), interpolation=InterpolationMode.BICUBIC), T.Pad(10), T.RandomCrop((args.height, args.width)), T.RandomHorizontalFlip(0.5), T.RandomRotation(5), T.ColorJitter(brightness=(0.5, 2.0), contrast=(0.5, 2.0), saturation=(0.5, 2.0), hue=(-0.1, 0.1)), T.RandomOcclusion(args.min_size, args.max_size), T.ToTensor(), ]) test_transformer = T.Compose([ T.Resize((args.height, args.width), interpolation=InterpolationMode.BICUBIC), T.ToTensor(), ]) train_path = osp.join(dataset.images_dir, dataset.train_path) train_loader = DataLoader( Preprocessor(dataset.train, root=train_path, transform=train_transformer), batch_size=args.batch_size, num_workers=args.workers, sampler=GraphSampler(dataset.train, train_path, test_transformer, model, matcher, args.batch_size, args.num_instance, args.test_gal_batch, args.test_prob_batch, save_path, args.gs_verbose), pin_memory=True) query_loader = DataLoader(Preprocessor(dataset.query, root=osp.join( dataset.images_dir, dataset.query_path), transform=test_transformer), batch_size=args.test_fea_batch, num_workers=args.workers, shuffle=False, pin_memory=True) gallery_loader = DataLoader(Preprocessor(dataset.gallery, root=osp.join( dataset.images_dir, dataset.gallery_path), transform=test_transformer), batch_size=args.test_fea_batch, num_workers=args.workers, shuffle=False, pin_memory=True) return dataset, num_classes, train_loader, query_loader, gallery_loader
def update_train_loader(dataset, train_samples, updated_label, height, width, batch_size, re, workers, all_img_cams, sample_position=7): normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) train_transformer = T.Compose([ T.Resize((height, width), interpolation=3), T.RandomHorizontalFlip(p=0.5), T.Pad(10), T.RandomCrop((height, width)), T.ToTensor(), normalizer, T.RandomErasing(EPSILON=re) ]) # obtain global accumulated label from pseudo label and cameras pure_label = updated_label[updated_label >= 0] pure_cams = all_img_cams[updated_label >= 0] accumulate_labels = np.zeros(pure_label.shape, pure_label.dtype) prev_id_count = 0 id_count_each_cam = [] for this_cam in np.unique(pure_cams): percam_labels = pure_label[pure_cams == this_cam] unique_id = np.unique(percam_labels) id_count_each_cam.append(len(unique_id)) id_dict = {ID: i for i, ID in enumerate(unique_id.tolist())} for i in range(len(percam_labels)): percam_labels[i] = id_dict[percam_labels[i]] accumulate_labels[pure_cams == this_cam] = percam_labels + prev_id_count prev_id_count += len(unique_id) print(' sum(id_count_each_cam)= {}'.format(sum(id_count_each_cam))) new_accum_labels = -1 * np.ones(updated_label.shape, updated_label.dtype) new_accum_labels[updated_label >= 0] = accumulate_labels # update sample list new_train_samples = [] for sample in train_samples: lbl = updated_label[sample[3]] if lbl != -1: assert (new_accum_labels[sample[3]] >= 0) new_sample = sample + (lbl, new_accum_labels[sample[3]]) new_train_samples.append(new_sample) target_train_loader = DataLoader(UnsupervisedTargetPreprocessor( new_train_samples, root=osp.join(dataset.target_images_dir, dataset.target_train_path), num_cam=dataset.target_num_cam, transform=train_transformer, has_pseudo_label=True), batch_size=batch_size, num_workers=workers, pin_memory=True, drop_last=True, sampler=ClassUniformlySampler( new_train_samples, class_position=sample_position, k=4)) return target_train_loader, len(new_train_samples)
def get_data(name, data_dir, height, width, batch_size, workers, combine_trainval, crop, tracking_icams, fps, re=0, num_instances=0, camstyle=0, zju=0, colorjitter=0): # if name == 'market1501': # root = osp.join(data_dir, 'Market-1501-v15.09.15') # elif name == 'duke_reid': # root = osp.join(data_dir, 'DukeMTMC-reID') # elif name == 'duke_tracking': # root = osp.join(data_dir, 'DukeMTMC') # else: # root = osp.join(data_dir, name) if name == 'duke_tracking': if tracking_icams != 0: tracking_icams = [tracking_icams] else: tracking_icams = list(range(1, 9)) dataset = datasets.create(name, data_dir, data_type='tracking_gt', iCams=tracking_icams, fps=fps, trainval=combine_trainval) elif name == 'aic_tracking': dataset = datasets.create(name, data_dir, data_type='tracking_gt', fps=fps, trainval=combine_trainval) else: dataset = datasets.create(name, data_dir) normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) num_classes = dataset.num_train_ids train_transformer = T.Compose([ T.ColorJitter(brightness=0.1 * colorjitter, contrast=0.1 * colorjitter, saturation=0.1 * colorjitter, hue=0), T.Resize((height, width)), T.RandomHorizontalFlip(), T.Pad(10 * crop), T.RandomCrop((height, width)), T.ToTensor(), normalizer, T.RandomErasing(probability=re), ]) test_transformer = T.Compose([ T.Resize((height, width)), # T.RectScale(height, width, interpolation=3), T.ToTensor(), normalizer, ]) if zju: train_loader = DataLoader( Preprocessor(dataset.train, root=dataset.train_path, transform=train_transformer), batch_size=batch_size, num_workers=workers, sampler=ZJU_RandomIdentitySampler(dataset.train, batch_size, num_instances) if num_instances else None, shuffle=False if num_instances else True, pin_memory=True, drop_last=False if num_instances else True) else: train_loader = DataLoader( Preprocessor(dataset.train, root=dataset.train_path, transform=train_transformer), batch_size=batch_size, num_workers=workers, sampler=RandomIdentitySampler(dataset.train, num_instances) if num_instances else None, shuffle=False if num_instances else True, pin_memory=True, drop_last=True) query_loader = DataLoader( Preprocessor(dataset.query, root=dataset.query_path, transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) gallery_loader = DataLoader( Preprocessor(dataset.gallery, root=dataset.gallery_path, transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) if camstyle <= 0: camstyle_loader = None else: camstyle_loader = DataLoader( Preprocessor(dataset.camstyle, root=dataset.camstyle_path, transform=train_transformer), batch_size=camstyle, num_workers=workers, shuffle=True, pin_memory=True, drop_last=True) return dataset, num_classes, train_loader, query_loader, gallery_loader, camstyle_loader
def get_data(data_dir, source, target, height, width, batch_size, re=0, workers=8): dataset = DA(data_dir, source, target) dataset_2 = DA(data_dir, target, source) normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) num_classes = dataset.num_train_ids train_transformer = T.Compose([ T.Resize((height, width), interpolation=3), T.RandomHorizontalFlip(p=0.5), T.Pad(10), T.RandomCrop((height, width)), T.ToTensor(), normalizer, # T.RandomErasing(EPSILON=re), T.RandomErasing(probability=0.4, mean=[0.485, 0.456, 0.406]) ]) train_transformer_2 = T.Compose([ T.Resize((height, width), interpolation=3), T.RandomHorizontalFlip(p=0.5), T.Pad(10), T.RandomCrop((height, width)), T.ToTensor(), normalizer, # T.RandomErasing(EPSILON=re), T.RandomErasing(probability=0.4, mean=[0.485, 0.456, 0.406]) ]) test_transformer = T.Compose([ T.Resize((height, width), interpolation=3), T.ToTensor(), normalizer, ]) ''' num_instances=4 rmgs_flag = num_instances > 0 if rmgs_flag: sampler = RandomIdentitySampler(dataset.target_train, num_instances) else: sampler = None ''' source_train_loader = DataLoader( Preprocessor(dataset.source_train, root=osp.join(dataset.source_images_dir, dataset.source_train_path), transform=train_transformer), batch_size=batch_size, num_workers=workers, shuffle=True, pin_memory=True, drop_last=True) num_instances=0 rmgs_flag = num_instances > 0 if rmgs_flag: sampler = RandomIdentitySampler(dataset.target_train, num_instances) else: sampler = None target_train_loader = DataLoader( UnsupervisedCamStylePreprocessor(dataset.target_train, root=osp.join(dataset.target_images_dir, dataset.target_train_path), camstyle_root=osp.join(dataset.target_images_dir, dataset.target_train_camstyle_path), num_cam=dataset.target_num_cam, transform=train_transformer), batch_size=batch_size, num_workers=workers, shuffle=True, pin_memory=True, drop_last=True) query_loader = DataLoader( Preprocessor(dataset.query, root=osp.join(dataset.target_images_dir, dataset.query_path), transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) query_loader_2 = DataLoader( Preprocessor(dataset_2.query, root=osp.join(dataset_2.target_images_dir, dataset_2.query_path), transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) gallery_loader = DataLoader( Preprocessor(dataset.gallery, root=osp.join(dataset.target_images_dir, dataset.gallery_path), transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) gallery_loader_2 = DataLoader( Preprocessor(dataset_2.gallery, root=osp.join(dataset_2.target_images_dir, dataset_2.gallery_path), transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) return dataset,dataset_2, num_classes, source_train_loader, target_train_loader, query_loader, gallery_loader, query_loader_2, gallery_loader_2
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 data loaders dataset, num_classes, train_loader, _, _ = \ get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.combine_all, args.workers, args.test_fea_batch) # Create model #model = seTest.resnst50().cuda() model = resmap.create(args.arch, ibn_type=args.ibn, final_layer=args.final_layer, neck=args.neck).cuda() num_features = model.num_features #num_features = 64 # 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 = ClassMemoryLoss(matcher, num_classes, num_features, hei, wid, args.mem_batch_size).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': criterion.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)) elif args.pre_epochs > 0: pre_tr = PreTrainer(model, criterion, optimizer, train_loader, args.pre_epochs, args.max_steps, args.num_trials) result_file = osp.join(exp_database_dir, args.method, 'pretrain_metric.txt') model, criterion, optimizer = pre_tr.train(result_file, args.method, args.sub_method) # 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) model = nn.DataParallel(model).cuda() criterion = nn.DataParallel(criterion).cuda() enhance_data_aug = False if not args.evaluate: # Trainer trainer = Trainer(model, criterion) 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.module.state_dict(), 'optim': optimizer.state_dict(), 'epoch': epoch + 1, }, fpath=osp.join(output_dir, 'checkpoint.pth.tar')) if not enhance_data_aug and epoch < args.epochs - 1 and acc > args.acc_thr: enhance_data_aug = True print('\nAcc = %.2f%% > %.2f%%. Start to Flip and Block.\n' % (acc * 100, args.acc_thr * 100)) train_transformer = T.Compose([ T.Resize((args.height, args.width), interpolation=3), T.Pad(10), T.RandomCrop((args.height, args.width)), T.RandomHorizontalFlip(0.5), T.RandomRotation(5), T.ColorJitter(brightness=(0.5, 2.0), contrast=(0.5, 2.0), saturation=(0.5, 2.0), hue=(-0.1, 0.1)), T.RandomOcclusion(args.min_size, args.max_size), T.ToTensor(), ]) train_loader = DataLoader(Preprocessor( dataset.train, root=osp.join(dataset.images_dir, dataset.train_path), transform=train_transformer), batch_size=args.batch_size, num_workers=args.workers, shuffle=True, pin_memory=True, drop_last=True) # 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) # With the seed reset (every time), the same set of numbers will appear every time. torch.manual_seed(args.seed) # Sets the seed for generating random numbers. cudnn.benchmark = True # This flag allows you to enable the inbuilt cudnn auto-tuner to find the best algorithm to use for your hardware. It enables benchmark mode in cudnn. # Redirect print to both console and log file sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt')) print(args) shutil.copy(sys.argv[0], osp.join(args.logs_dir, osp.basename(sys.argv[0]))) # Create data loaders 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 * 8, args.workers, #https://deeplizard.com/learn/video/kWVgvsejXsE#:~:text=The%20num_workers%20attribute%20tells%20the,sequentially%20inside%20the%20main%20process. ) # Create model model = models.create("ft_net_inter", num_classes=num_classes, stride=args.stride) # Load from checkpoint start_epoch = 0 best_top1 = 0 top1 = 0 is_best = False if args.checkpoint is not None: if args.evaluate: checkpoint = load_checkpoint(args.checkpoint) param_dict = model.state_dict() # A state_dict is simply a Python dictionary object that maps each layer to its parameter tensor. for k, v in checkpoint['state_dict'].items(): if 'model' in k: param_dict[k] = v model.load_state_dict(param_dict) else: model.model.load_param(args.checkpoint) model = model.cuda() # Distance metric metric = None # Evaluator evaluator = Evaluator(model, use_cpu=args.use_cpu) if args.evaluate: print("Test:") evaluator.evaluate(test_loader, dataset.query, dataset.gallery, metric) return train_transformer = [ T.Resize((args.height, args.width), interpolation=3), T.RandomHorizontalFlip(), T.Pad(10), T.RandomCrop((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), T.RandomErasing(probability=0.5), ] train_transformer = T.Compose(train_transformer) for cluster_epoch in range(args.cluster_epochs): # -------------------------Stage 1 intra camera training-------------------------- # Cluster and generate new dataset and model # Divides the training set into (subsets) and according to that each camera id is there for each image # then it forms clustering on each subset according to the pair wise similarity # then assigning images with in each cluster with identical label # then cross entropy loss is used cluster_result = get_intra_cam_cluster_result(model, train_loader, args.class_number_stage1, args.linkage) cluster_datasets = [ datasets.create("cluster", osp.join(args.data_dir, args.dataset), cluster_result[cam_id], cam_id) for cam_id in cluster_result.keys() ] cluster_dataloaders = [ DataLoader(Preprocessor(dataset.train_set, root=dataset.images_dir, transform=train_transformer), batch_size=args.batch_size, num_workers=args.workers, shuffle=True, pin_memory=False, drop_last=True) for dataset in cluster_datasets ] param_dict = model.model.state_dict() model = models.create("ft_net_intra", num_classes=[ args.class_number_stage1 for cam_id in cluster_result.keys() ], stride=args.stride) model_param_dict = model.model.state_dict() for k, v in model_param_dict.items(): if k in param_dict.keys(): model_param_dict[k] = param_dict[k] model.model.load_state_dict(model_param_dict) model = model.cuda() criterion = nn.CrossEntropyLoss().cuda() # Optimizer param_groups = make_params(model, args.lr, args.weight_decay) optimizer = torch.optim.SGD(param_groups, momentum=0.9) # Trainer trainer = IntraCameraTrainer( model, criterion, warm_up_epoch=args.warm_up) print("start training") # Start training for epoch in range(0, args.epochs_stage1): trainer.train(cluster_epoch, epoch, cluster_dataloaders, optimizer, print_freq=args.print_freq, ) # -------------------------------------------Stage 2 inter camera training----------------------------------- mix_rate = get_mix_rate( args.mix_rate, cluster_epoch, args.cluster_epochs, power=args.decay_factor) cluster_result = get_inter_cam_cluster_result( model, train_loader, args.class_number_stage2, args.linkage, mix_rate, use_cpu=args.use_cpu) cluster_dataset = datasets.create( "cluster", osp.join(args.data_dir, args.dataset), cluster_result, 0) cluster_dataloaders = DataLoader( Preprocessor(cluster_dataset.train_set, root=cluster_dataset.images_dir, transform=train_transformer), batch_size=args.batch_size_stage2, num_workers=args.workers, sampler=RandomIdentitySampler(cluster_dataset.train_set, args.batch_size_stage2, args.instances), pin_memory=False, drop_last=True) param_dict = model.model.state_dict() model = models.create("ft_net_inter", num_classes=args.class_number_stage2, stride=args.stride) model.model.load_state_dict(param_dict) model = model.cuda() # Criterion criterion_entropy = nn.CrossEntropyLoss().cuda() criterion_triple = TripletLoss(margin=args.margin).cuda() # Optimizer param_groups = make_params(model, args.lr * args.batch_size_stage2 / 32, args.weight_decay) optimizer = torch.optim.SGD(param_groups, momentum=0.9) # Trainer trainer = InterCameraTrainer(model, criterion_entropy, criterion_triple, warm_up_epoch=args.warm_up, ) print("start training") # Start training for epoch in range(0, args.epochs_stage2): trainer.train(cluster_epoch, epoch, cluster_dataloaders, optimizer, print_freq=args.print_freq) if (cluster_epoch + 1) % 5 == 0: # in 4th, 9th, 14th epochs, see in the output evaluator = Evaluator(model, use_cpu=args.use_cpu) top1, mAP = evaluator.evaluate( test_loader, dataset.query, dataset.gallery, metric, return_mAP=True) is_best = top1 > best_top1 best_top1 = max(top1, best_top1) save_checkpoint( { 'state_dict': model.state_dict(), 'epoch': cluster_epoch + 1, 'best_top1': best_top1, 'cluster_epoch': cluster_epoch + 1, }, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar')) if cluster_epoch == (args.cluster_epochs - 1): save_checkpoint( { 'state_dict': model.state_dict(), 'epoch': cluster_epoch + 1, 'best_top1': best_top1, 'cluster_epoch': cluster_epoch + 1, }, False, fpath=osp.join(args.logs_dir, 'latest.pth.tar')) print('\n * cluster_epoch: {:3d} top1: {:5.1%} best: {:5.1%}{}\n'. format(cluster_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']) best_rank1, mAP = evaluator.evaluate( test_loader, dataset.query, dataset.gallery, metric, return_mAP=True)