def transform(self): transform_train = [ transforms.Resize((256, 128), interpolation=3), transforms.Pad(10), transforms.RandomCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] transform_val = [ transforms.Resize(size=(256, 128), interpolation=3), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if self.erasing_p > 0: transform_train = transform_train + [ RandomErasing(probability=self.erasing_p, mean=[0.0, 0.0, 0.0]) ] if self.color_jitter: transform_train = [ transforms.ColorJitter( brightness=0.1, contrast=0.1, saturation=0.1, hue=0) ] + transform_train self.data_transforms = { 'train': transforms.Compose(transform_train), 'val': transforms.Compose(transform_val), }
def __init__(self, args): train_list = [ transforms.Resize((args.height, args.width), interpolation=3), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ] if args.random_erasing: train_list.append( RandomErasing(probability=args.probability, mean=[0.0, 0.0, 0.0])) train_transform = transforms.Compose(train_list) test_transform = transforms.Compose([ transforms.Resize((args.height, args.width), interpolation=3), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) if not args.test_only: module_train = import_module('data.' + args.data_train.lower()) self.trainset = getattr(module_train, args.data_train)(args, train_transform, 'train') self.train_loader = dataloader.DataLoader( self.trainset, sampler=RandomSampler(self.trainset, args.batchid, batch_image=args.batchimage), #shuffle=True, batch_size=args.batchid * args.batchimage, num_workers=args.nThread) else: self.train_loader = None if args.data_test in ['Market1501']: module = import_module('data.' + args.data_train.lower()) self.testset = getattr(module, args.data_test)(args, test_transform, 'test') self.queryset = getattr(module, args.data_test)(args, test_transform, 'query') else: raise Exception() self.test_loader = dataloader.DataLoader(self.testset, batch_size=args.batchtest, num_workers=args.nThread) self.query_loader = dataloader.DataLoader(self.queryset, batch_size=args.batchtest, num_workers=args.nThread)
def __init__(self, data_cfg, multi=1, nl=False): """ Dataset for training. :param data_cfg: CfgNode for CityFlow NL. """ self.nl = nl self.multi = multi self.motion = data_cfg.motion self.nseg = data_cfg.nseg self.all3 = data_cfg.all3 self.pad = data_cfg.pad self.data_cfg = data_cfg self.aug = AutoAugment(auto_augment_policy(name='v0r', hparams=None)) with open(self.data_cfg.JSON_PATH) as f: tracks = json.load(f) f.close() self.list_of_uuids = list(tracks.keys()) self.list_of_tracks = list(tracks.values()) self.list_of_crops = list() train_num = len(self.list_of_uuids) self.transform = transforms.Compose([ transforms.Pad(10), transforms.RandomCrop( (data_cfg.CROP_SIZE, self.data_cfg.CROP_SIZE)), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), RandomErasing(probability=0.5) ]) if data_cfg.semi: #cv with open(self.data_cfg.EVAL_TRACKS_JSON_PATH) as f: unlabel_tracks = json.load(f) f.close() self.list_of_uuids.extend(unlabel_tracks.keys()) self.list_of_tracks.extend(unlabel_tracks.values()) #nl with open("data/test-queries.json", "r") as f: unlabel_nl = json.load(f) unlabel_nl_key = list(unlabel_nl.keys()) print('#track id (class): %d ' % len(self.list_of_tracks)) count = 0 # add id and nl, -1 for unlabeled data for track_idx, track in enumerate(self.list_of_tracks): track["track_id"] = track_idx track["nl_id"] = track_idx # from 0 to train_num-1 is the id of the original training set. if track_idx >= train_num: track["nl_id"] = -1 track["nl"] = unlabel_nl[unlabel_nl_key[count]] count = count + 1 self._logger = get_logger()
def __init__(self, data_pth, is_train=True, *args, **kwargs): super(Market1501, self).__init__(*args, **kwargs) ## parse image names to generate image ids imgs = os.listdir(data_pth) imgs = [im for im in imgs if osp.splitext(im)[-1] == '.jpg'] self.is_train = is_train self.im_pths = [osp.join(data_pth, im) for im in imgs] self.im_infos = {} self.person_infos = {} for i, im in enumerate(imgs): tokens = im.split('_') im_pth = self.im_pths[i] pid = int(tokens[0]) cam = int(tokens[1][1]) self.im_infos.update({im_pth: (pid, cam)}) if pid in self.person_infos.keys(): self.person_infos[pid].append(i) else: self.person_infos[pid] = [ i, ] self.pid_label_map = {} for i, (pid, ids) in enumerate(self.person_infos.items()): self.person_infos[pid] = np.array(ids, dtype=np.int32) self.pid_label_map[pid] = i ## preprocessing self.trans_train = transforms.Compose([ transforms.Resize((288, 144)), transforms.RandomCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), RandomErasing(0.5, mean=[0.0, 0.0, 0.0]) ]) ## H-Flip self.trans_no_train_flip = transforms.Compose([ transforms.Resize((288, 144)), transforms.RandomHorizontalFlip(1), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) self.trans_no_train_noflip = transforms.Compose([ transforms.Resize((288, 144)), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ])
def __init__(self, crop_prob=0, crop_ratio=1.0, resize_h_w=None, scale=True, im_mean=None, im_std=None, mirror_type=None, batch_dims='NCHW', is_random_erasing=False, prng=np.random): """ Args: crop_prob: the probability of each image to go through cropping crop_ratio: a float. If == 1.0, no cropping. resize_h_w: (height, width) after resizing. If `None`, no resizing. scale: whether to scale the pixel value by 1/255 im_mean: (Optionally) subtracting image mean; `None` or a tuple or list or numpy array with shape [3] im_std: (Optionally) divided by image std; `None` or a tuple or list or numpy array with shape [3]. Dividing is applied only when subtracting mean is applied. mirror_type: How image should be mirrored; one of [None, 'random', 'always'] batch_dims: either 'NCHW' or 'NHWC'. 'N': batch size, 'C': num channels, 'H': im height, 'W': im width. PyTorch uses 'NCHW', while TensorFlow uses 'NHWC'. prng: can be set to a numpy.random.RandomState object, in order to have random seed independent from the global one """ self.crop_prob = crop_prob self.crop_ratio = crop_ratio self.resize_h_w = resize_h_w self.scale = scale self.im_mean = im_mean self.im_std = im_std self.check_mirror_type(mirror_type) self.mirror_type = mirror_type self.check_batch_dims(batch_dims) self.batch_dims = batch_dims self.prng = prng self.is_random_erasing = is_random_erasing self.RandomErasing = RandomErasing(probability=0.5, sh=0.2, mean=[0.0, 0.0, 0.0])
def __init__(self, pids_list, array, is_train=True, *args, **kwargs): super(Market1501, self).__init__(*args, **kwargs) self.is_train = is_train self.pids = pids_list self.array = array self.img_names = [ el.split("/")[1] for el in self.pids if os.path.splitext(el)[1] == '.jpg' ] self.lb_ids = [int(el.split('_')[0]) for el in self.img_names] self.lb_cams = [int(el.split('_')[1][1]) for el in self.img_names] if is_train: self.trans = transforms.Compose([ transforms.ToPILImage(), transforms.Resize((288, 144)), transforms.RandomCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.486, 0.459, 0.408), (0.229, 0.224, 0.225)), RandomErasing(0.5, mean=[0.0, 0.0, 0.0]) ]) else: self.trans_tuple = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.486, 0.459, 0.408), (0.229, 0.224, 0.225)) ]) self.Lambda = transforms.Lambda( lambda crops: [self.trans_tuple(crop) for crop in crops]) self.trans = transforms.Compose([ transforms.ToPILImage(), transforms.Resize((288, 144)), transforms.TenCrop((256, 128)), self.Lambda ]) # useful for sampler self.lb_img_dict = dict() self.lb_ids_uniq = set(self.lb_ids) lb_array = np.array(self.lb_ids) for lb in self.lb_ids_uniq: idx = np.where(lb_array == lb)[0] self.lb_img_dict.update({lb: idx})
###################################################################### # Load Data # --------- # transform_train_list = [ #transforms.RandomResizedCrop(size=128, scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC) transforms.Resize(144, interpolation=3), transforms.RandomCrop((256,128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.erasing_p>0: transform_train_list = transform_train_list + [RandomErasing(opt.erasing_p)] if opt.color_jitter: transform_train_list = [transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0)] + transform_train_list print(transform_train_list) transform_val_list = [ transforms.Resize(size=(256,128),interpolation=3), #Image.BICUBIC transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] data_transforms = { 'train': transforms.Compose( transform_train_list ),
def main(): transforms_args = [ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] train_dataset = CoCoDataset( args.coco_path, "training", target_size=args.target_size, transform=transforms.Compose( transforms_args + [RandomErasing(probability=args.p, sh=args.sh, r1=args.r1)])) test_dataset = CoCoDataset(args.coco_path, "validation_wo_occlusion", target_size=args.target_size, transform=transforms.Compose(transforms_args)) train_batch_sampler = TrainBalancedBatchSampler(torch.from_numpy( np.array(train_dataset.all_targets())), K=args.K, P=args.P, n_batches=args.n_batches) test_batch_sampler = TestBalancedBatchSampler(torch.from_numpy( np.array(test_dataset.all_targets())), K=args.K, P=args.P, n_batches=args.n_batches) train_loader = DataLoader(train_dataset, batch_sampler=train_batch_sampler, **kwargs) test_loader = DataLoader(test_dataset, batch_sampler=test_batch_sampler, **kwargs) # init model model, optim_state_dict, init_epoch = load_model( args.backbone, args.snapshot, imagenet_weights=args.imagenet_weights, freeze=args.freeze) print("Resume training from epoch", init_epoch) if cuda: model.cuda() # init optimizer if args.optim == "Adam": optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) elif args.optim == "SGD": optimizer = optim.SGD(model.parameters(), momentum=0.9, lr=args.lr, weight_decay=1e-4) else: raise ValueError("Optimizer is not supported") if optim_state_dict is not None: optimizer.load_state_dict(optim_state_dict) # define loss function if args.triplet_selector == "hard": selector = HardestNegativeTripletSelector(args.soft_margin) elif args.triplet_selector == "semi": selector = SemihardNegativeTripletSelector(args.soft_margin) elif args.triplet_selector == "random": selector = RandomNegativeTripletSelector(args.soft_margin) else: selector = AllTripletSelector() train_loss_fn = TripletLoss(selector, soft_margin=args.soft_margin) test_loss_fn = TripletLoss(AllTripletSelector(), soft_margin=args.soft_margin) # define learning rate scheduler lr_scheduler = LrScheduler(args.epoch_decay_start, args.n_epoch, args.lr) log_file = os.path.join( args.logger_dir, '%s_%s.csv' % (args.backbone, args.triplet_selector)) for epoch in range(init_epoch + 1, args.n_epoch): lr_scheduler.adjust_learning_rate(optimizer, epoch, args.optim) for param_group in optimizer.param_groups: print("LR: ", param_group['lr']) train_loss = train_epoch(model, train_loader, train_loss_fn, optimizer, cuda) if epoch % args.eval_freq == 0: test_loss = test_epoch(model, test_loader, test_loss_fn, cuda) print('Epoch [%d/%d], Train loss: %.4f, Test loss: %.4f' % (epoch, args.n_epoch, train_loss, test_loss)) log = [epoch, train_loss, test_loss] if os.path.isfile(log_file): with open(log_file, mode='a', newline='') as csv_f: writer = csv.writer(csv_f) writer.writerow(log) else: with open(log_file, mode='w', newline='') as csv_f: writer = csv.writer(csv_f) # write header writer.writerow(["epoch", "train_loss", "test_loss"]) writer.writerow(log) if epoch % args.save_freq == 0: torch.save( { 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'epoch': epoch }, os.path.join( args.snapshot_path, '%s_%s_%d.pth' % (args.backbone, args.triplet_selector, epoch)))
transforms.RandomHorizontalFlip(0.5), transforms.Pad(10), transforms.RandomCrop([288, 144]), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] transform_val_list = [ transforms.Resize(size=(256, 128), interpolation=3), #Image.BICUBIC transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if erasing_p > 0: transform_train_list = transform_train_list + [ RandomErasing(probability=erasing_p) ] print(transform_train_list) data_transforms = { 'train': transforms.Compose(transform_train_list), 'val': transforms.Compose(transform_val_list), } train_all = '' if train_all_1: train_all = '_all' image_datasets = {} image_datasets['train'] = datasets.ImageFolder( os.path.join(data_dir, 'train' + train_all), data_transforms['train'])
def main(): global best_rank1 global pn_gan global poses if args.use_gan: pn_gan, poses = initialize_PN_GAN() torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset( root=args.root, name=args.dataset, split_id=args.split_id, cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split, ) transform_train_list = ([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) if args.erasing_p > 0: transform_train_list = transform_train_list + [ RandomErasing(probability=args.erasing_p, mean=[0.0, 0.0, 0.0]) ] transform_train = T.Compose(transform_train_list) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( ImageDataset(dataset.query, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( ImageDataset(dataset.gallery, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent'}, use_gpu=use_gpu) print("Model size: {:.3f} M".format(count_num_param(model))) if args.label_smooth: criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) else: criterion = nn.CrossEntropyLoss() optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.fixbase_epoch > 0: if hasattr(model, 'classifier') and isinstance(model.classifier, nn.Module): optimizer_tmp = init_optim(args.optim, model.classifier.parameters(), args.fixbase_lr, args.weight_decay) else: print( "Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0" ) args.fixbase_epoch = 0 if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size model = torch.load('') model.eval() print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint_saved = torch.load('./checkpoint_self_after_reset8.pth.tar') model.load_state_dict(checkpoint_saved) model.train() print("Loaded checkpoint from '{}'".format(args.resume)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results( distmat, dataset, save_dir=osp.join(args.save_dir, 'ranked_results'), topk=20, ) return start_time = time.time() train_time = 0 best_epoch = args.start_epoch print("==> Start training") if args.fixbase_epoch > 0: print( "Train classifier for {} epochs while keeping base network frozen". format(args.fixbase_epoch)) for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer_tmp, trainloader, use_gpu, freeze_bn=True) train_time += round(time.time() - start_train_time) del optimizer_tmp print("Now open all layers for training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() torch.save(state_dict, 'checkpoint_self_after_reset' + str(epoch + 1) + '.pth.tar') if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, queryloader, galleryloader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def __init__(self, data_path, train_file, img_size, is_train=True, *args, **kwargs): super(VehicleID, self).__init__(*args, **kwargs) self.is_train = is_train self.data_path = data_path self.check = os.listdir(data_path) self.check = [ el for el in self.check if os.path.splitext(el)[1] == '.jpg' ] self.dict = {} f = open(train_file, 'r') self.train_file = f.readlines() self.imgs = [] max_class = 0 self.classes = set() now = 0 for line in self.train_file: line = line.strip() if (len(line.split(' ')) == 2): car_name, car_class = line.split(' ')[0], int( line.split(' ')[1]) elif (len(line.split(' ')) == 1 and not is_train): car_name, car_class = line.split(' ')[0], -1 else: logger.info('dataset wrong') #if car_name not in self.check: # logger.warning("%s do not exists"%(car_name)) # continue if (len(car_name.split('.')) == 1): car_name = car_name + '.jpg' if car_name not in self.dict: self.dict[car_name] = car_class max_class = max(max_class, int(car_class)) self.classes.add(car_class) self.imgs.append(car_name) logger.info('dataset OK') if self.is_train: assert (max_class == len(self.classes) - 1) self.lb_ids = [self.dict[el] for el in self.imgs] self.lb_cams = [-1 for el in self.imgs] self.imgs = [os.path.join(data_path, el) for el in self.imgs] if is_train: self.trans = transforms.Compose([ transforms.Resize((img_size, img_size)), transforms.RandomCrop((img_size, img_size)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.486, 0.459, 0.408), (0.229, 0.224, 0.225)), RandomErasing(0.5, mean=[0.0, 0.0, 0.0]) ]) else: self.trans_tuple = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.486, 0.459, 0.408), (0.229, 0.224, 0.225)) ]) self.Lambda = transforms.Lambda( lambda crops: [self.trans_tuple(crop) for crop in crops]) self.trans = transforms.Compose([ transforms.Resize((img_size, img_size)), transforms.TenCrop((img_size, img_size)), self.Lambda, ]) # useful for sampler self.lb_img_dict = dict() self.lb_ids_uniq = set(self.lb_ids) lb_array = np.array(self.lb_ids) for lb in self.lb_ids_uniq: idx = np.where(lb_array == lb)[0] self.lb_img_dict.update({lb: idx})
from train_config import args, params grad_list = [] def print_grad(grad): print(grad) grad_list.append(grad) data_transforms = { 'train': transforms.Compose([ transforms.Resize((224,224)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.1,0.1,0.1,0), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), RandomErasing(probability=0.5, mean=[0.0, 0.0, 0.0]) ]), 'val': transforms.Compose([ transforms.Resize((224,224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } train_datasets = PETAData(txt_path='data_list/PETA_RAP_train.txt',dataset='train', data_transforms = data_transforms) test_datasets = PETAData(txt_path='data_list/PETA_RAP_test.txt', dataset='val', data_transforms = data_transforms) image_datasets = {'train' : train_datasets, 'val' : test_datasets}
def main(): print("==========\nArgs:{}\n==========".format(opt)) train_data_dir = opt.train_data_dir test_data_dir = opt.test_data_dir name = opt.name os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpu_ids use_gpu = torch.cuda.is_available() if opt.use_cpu: use_gpu = False if not opt.evaluate: sys.stdout = Logger(osp.join(opt.save_dir, opt.train_log)) else: sys.stdout = Logger(osp.join(opt.save_dir, opt.test_log)) if use_gpu: print("Currently using GPU {}".format(opt.gpu_ids)) cudnn.benchmark = True else: print("Currently using CPU (GPU is highly recommended)") #str_ids = opt.gpu_ids.split(',') #gpu_ids = [] #for str_id in str_ids: # gid = int(str_id) # if gid >=0: # gpu_ids.append(gid) # ## set gpu ids #if len(gpu_ids)>0: # torch.cuda.set_device(gpu_ids[0]) #print(gpu_ids[0]) # Load train Data # --------- print("==========Preparing trian dataset========") transform_train_list = [ # transforms.RandomResizedCrop(size=128, scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC) transforms.Resize((288, 144), interpolation=3), transforms.RandomCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.PCB: transform_train_list = [ transforms.Resize((384, 192), interpolation=3), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.erasing_p > 0: transform_train_list = transform_train_list + [ RandomErasing(probability=opt.erasing_p, mean=[0.0, 0.0, 0.0]) ] if opt.color_jitter: transform_train_list = [ transforms.ColorJitter( brightness=0.1, contrast=0.1, saturation=0.1, hue=0) ] + transform_train_list train_data_transforms = transforms.Compose(transform_train_list) train_all = '' if opt.train_all: if opt.use_clean_imgs: train_all = '_all_clean' else: train_all = '_all' print("Using all the train images") train_image_datasets = {} train_image_datasets['train'] = datasets.ImageFolder( os.path.join(train_data_dir, 'train' + train_all), train_data_transforms) train_dataloaders = { x: torch.utils.data.DataLoader(train_image_datasets[x], batch_size=opt.train_batch, shuffle=True, num_workers=4) for x in ['train'] } dataset_sizes = {x: len(train_image_datasets[x]) for x in ['train']} class_names = train_image_datasets['train'].classes inputs, classes = next(iter(train_dataloaders['train'])) ###################################################################### # Prepare test data if not opt.train_only: print("========Preparing test dataset========") transform_test_list = transforms.Compose([ transforms.Resize((384, 192), interpolation=3), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) test_image_datasets = { x: datasets.ImageFolder(os.path.join(test_data_dir, x), transform_test_list) for x in ['gallery', 'query'] } test_dataloaders = { x: torch.utils.data.DataLoader(test_image_datasets[x], batch_size=opt.test_batch, shuffle=False, num_workers=4) for x in ['gallery', 'query'] } print("Initializing model...") if opt.use_dense: model = ft_net_dense(len(class_names)) else: model = ft_net(len(class_names)) if opt.PCB: model = PCB(len(class_names)) print("Model size: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) start_epoch = opt.start_epoch if opt.resume: print("Loading checkpoint from '{}'".format(opt.resume)) checkpoint = torch.load(opt.resume) model.load_state_dict(checkpoint['state_dict']) # model.load_state_dict(checkpoint) start_epoch = checkpoint['epoch'] if use_gpu: model = nn.DataParallel(model).cuda() if opt.evaluate: print("Evaluate only") test(model, test_image_datasets, test_dataloaders, use_gpu) return criterion = nn.CrossEntropyLoss().cuda() if opt.PCB: ignored_params = list(map(id, model.module.resnet50.fc.parameters())) ignored_params += ( list(map(id, model.module.classifier0.parameters())) + list(map(id, model.module.classifier1.parameters())) + list(map(id, model.module.classifier2.parameters())) + list(map(id, model.module.classifier3.parameters())) + list(map(id, model.module.classifier4.parameters())) + list(map(id, model.module.classifier5.parameters())) #+list(map(id, model.classifier7.parameters() )) ) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer = optim.SGD( [ { 'params': base_params, 'lr': 0.01 }, { 'params': model.module.resnet50.fc.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier0.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier1.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier2.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier3.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier4.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier5.parameters(), 'lr': 0.1 }, #{'params': model.classifier7.parameters(), 'lr': 0.01} ], weight_decay=5e-4, momentum=0.9, nesterov=True) else: ignored_params = list(map( id, model.module.model.fc.parameters())) + list( map(id, model.module.classifier.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer = optim.SGD([{ 'params': base_params, 'lr': 0.01 }, { 'params': model.module.model.fc.parameters(), 'lr': 0.1 }, { 'params': model.module.classifier.parameters(), 'lr': 0.1 }], weight_decay=5e-4, momentum=0.9, nesterov=True) # Decay LR by a factor of 0.1 every 40 epochs if opt.stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.stepsize, gamma=0.1) start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") ###################################################################### # Training the model # ------------------ # # Now, let's write a general function to train a model. Here, we will # illustrate: # # - Scheduling the learning rate # - Saving the best model # # In the following, parameter ``scheduler`` is an LR scheduler object from # ``torch.optim.lr_scheduler``. for epoch in range(start_epoch, opt.max_epoch): start_train_time = time.time() if opt.train_only: print("==> Training only") train(epoch, model, criterion, optimizer, train_dataloaders, use_gpu) train_time += round(time.time() - start_train_time) if epoch % opt.eval_step == 0: if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'epoch': epoch, }, 0, osp.join(opt.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) if opt.stepsize > 0: scheduler.step() else: train(epoch, model, criterion, optimizer, train_dataloaders, use_gpu) train_time += round(time.time() - start_train_time) if opt.stepsize > 0: scheduler.step() if (epoch + 1) > opt.start_eval and opt.eval_step > 0 and ( epoch + 1) % opt.eval_step == 0 or (epoch + 1) == opt.max_epoch: print("==> Test") rank1 = test(model, test_image_datasets, test_dataloaders, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(opt.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) if not opt.train_only: print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def __init__(self, **kwargs): super(RandomErasingImageAugmentor, self).__init__(**kwargs) additional_transforms = [RandomErasing()] self.add_transforms(additional_transforms)
def __init__(self, hyperparameters, gpu_ids=[0]): super(DGNet_Trainer, self).__init__() # 从配置文件获取生成模型和鉴别模型的学习率 lr_g = hyperparameters['lr_g'] lr_d = hyperparameters['lr_d'] # ID 类别 ID_class = hyperparameters['ID_class'] # 是否设置使用fp16, if not 'apex' in hyperparameters.keys(): hyperparameters['apex'] = False self.fp16 = hyperparameters['apex'] # Initiate the networks # We do not need to manually set fp16 in the network for the new apex. So here I set fp16=False. self.gen_a = AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen'], fp16=False) # auto-encoder for domain a self.gen_b = self.gen_a # auto-encoder for domain b ''' ft_netAB : Ea ''' # ID_stride: 外观编码器池化层的stride if not 'ID_stride' in hyperparameters.keys(): hyperparameters['ID_stride'] = 2 # id_a : 外观编码器 -> Ea if hyperparameters['ID_style'] == 'PCB': self.id_a = PCB(ID_class) elif hyperparameters['ID_style'] == 'AB': self.id_a = ft_netAB(ID_class, stride=hyperparameters['ID_stride'], norm=hyperparameters['norm_id'], pool=hyperparameters['pool']) else: self.id_a = ft_net(ID_class, norm=hyperparameters['norm_id'], pool=hyperparameters['pool']) # return 2048 now self.id_b = self.id_a # 对图片b的操作与图片a的操作一致 # 判别器,使用的是一个多尺寸的判别器,就是对图片进行几次缩放,并且对每次缩放都会预测,计算总的损失 # 经过网络3个缩放,,分别为:[batch_size, 1, 64, 32],[batch_size, 1, 32, 16],[batch_size, 1, 16, 8] self.dis_a = MsImageDis(3, hyperparameters['dis'], fp16=False) # discriminator for domain a self.dis_b = self.dis_a # discriminator for domain b # load teachers if hyperparameters['teacher'] != "": teacher_name = hyperparameters['teacher'] print(teacher_name) # 加载多个老师模型 teacher_names = teacher_name.split(',') # 构建老师模型 teacher_model = nn.ModuleList() # 初始化为空,接下来开始填充 teacher_count = 0 for teacher_name in teacher_names: config_tmp = load_config(teacher_name) # 池化层的stride if 'stride' in config_tmp: stride = config_tmp['stride'] else: stride = 2 # 开始搭建网络 model_tmp = ft_net(ID_class, stride=stride) teacher_model_tmp = load_network(model_tmp, teacher_name) teacher_model_tmp.model.fc = nn.Sequential( ) # remove the original fc layer in ImageNet teacher_model_tmp = teacher_model_tmp.cuda() # teacher_model_tmp,[3, 224, 224] # 使用fp16 if self.fp16: teacher_model_tmp = amp.initialize(teacher_model_tmp, opt_level="O1") teacher_model.append(teacher_model_tmp.cuda().eval( )) # 第一个填充为 teacher_model_tmp.cuda().eval() teacher_count += 1 self.teacher_model = teacher_model # 是否使用batchnorm if hyperparameters['train_bn']: self.teacher_model = self.teacher_model.apply(train_bn) # 实例正则化 self.instancenorm = nn.InstanceNorm2d(512, affine=False) # RGB to one channel # 因为Es 需要使用灰度图, 所以single 用来将图片转化为灰度图 if hyperparameters['single'] == 'edge': self.single = to_edge else: self.single = to_gray(False) # Random Erasing when training # arasing_p 随机擦除的概率 if not 'erasing_p' in hyperparameters.keys(): self.erasing_p = 0 else: self.erasing_p = hyperparameters['erasing_p'] # 对图片中的某一随机区域进行擦除,具体:将该区域的像素值设置为均值 self.single_re = RandomErasing(probability=self.erasing_p, mean=[0.0, 0.0, 0.0]) if not 'T_w' in hyperparameters.keys(): hyperparameters['T_w'] = 1 # Setup the optimizers beta1 = hyperparameters['beta1'] beta2 = hyperparameters['beta2'] dis_params = list( self.dis_a.parameters()) #+ list(self.dis_b.parameters()) gen_params = list( self.gen_a.parameters()) #+ list(self.gen_b.parameters()) self.dis_opt = torch.optim.Adam( [p for p in dis_params if p.requires_grad], lr=lr_d, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay']) self.gen_opt = torch.optim.Adam( [p for p in gen_params if p.requires_grad], lr=lr_g, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay']) # id params # 修改 id_a模型中分类器的学习率 if hyperparameters['ID_style'] == 'PCB': ignored_params = ( list(map(id, self.id_a.classifier0.parameters())) + list(map(id, self.id_a.classifier1.parameters())) + list(map(id, self.id_a.classifier2.parameters())) + list(map(id, self.id_a.classifier3.parameters()))) base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters()) lr2 = hyperparameters['lr2'] self.id_opt = torch.optim.SGD( [{ 'params': base_params, 'lr': lr2 }, { 'params': self.id_a.classifier0.parameters(), 'lr': lr2 * 10 }, { 'params': self.id_a.classifier1.parameters(), 'lr': lr2 * 10 }, { 'params': self.id_a.classifier2.parameters(), 'lr': lr2 * 10 }, { 'params': self.id_a.classifier3.parameters(), 'lr': lr2 * 10 }], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True) elif hyperparameters['ID_style'] == 'AB': ignored_params = ( list(map(id, self.id_a.classifier1.parameters())) + list(map(id, self.id_a.classifier2.parameters()))) base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters()) lr2 = hyperparameters['lr2'] self.id_opt = torch.optim.SGD( [{ 'params': base_params, 'lr': lr2 }, { 'params': self.id_a.classifier1.parameters(), 'lr': lr2 * 10 }, { 'params': self.id_a.classifier2.parameters(), 'lr': lr2 * 10 }], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True) else: ignored_params = list(map(id, self.id_a.classifier.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters()) lr2 = hyperparameters['lr2'] self.id_opt = torch.optim.SGD( [{ 'params': base_params, 'lr': lr2 }, { 'params': self.id_a.classifier.parameters(), 'lr': lr2 * 10 }], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True) # 生成器和判别器中的优化策略(学习率的更新策略) self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters) self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters) self.id_scheduler = get_scheduler(self.id_opt, hyperparameters) self.id_scheduler.gamma = hyperparameters['gamma2'] #ID Loss self.id_criterion = nn.CrossEntropyLoss() self.criterion_teacher = nn.KLDivLoss( size_average=False) # 生成主要特征: Lprim # Load VGG model if needed if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0: self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models') self.vgg.eval() for param in self.vgg.parameters(): param.requires_grad = False # save memory # 保存当前的模型,是为了提高计算效率 if self.fp16: # Name the FP16_Optimizer instance to replace the existing optimizer assert torch.backends.cudnn.enabled, "fp16 mode requires cudnn backend to be enabled." self.gen_a = self.gen_a.cuda() self.dis_a = self.dis_a.cuda() self.id_a = self.id_a.cuda() self.gen_b = self.gen_a self.dis_b = self.dis_a self.id_b = self.id_a self.gen_a, self.gen_opt = amp.initialize(self.gen_a, self.gen_opt, opt_level="O1") self.dis_a, self.dis_opt = amp.initialize(self.dis_a, self.dis_opt, opt_level="O1") self.id_a, self.id_opt = amp.initialize(self.id_a, self.id_opt, opt_level="O1")
def __init__(self, hyperparameters, gpu_ids=[0]): super(DGNet_Trainer, self).__init__() # 从配置文件获取生成模型的和鉴别模型的学习率 lr_g = hyperparameters['lr_g'] lr_d = hyperparameters['lr_d'] # # ID的类别,这里要注意,不同的数据集都是不一样的,应该是训练数据集的ID数目,非测试集 ID_class = hyperparameters['ID_class'] # 看是否设置使用float16,估计float16可以增加精确度 if not 'apex' in hyperparameters.keys(): hyperparameters['apex'] = False self.fp16 = hyperparameters['apex'] # Initiate the networks # We do not need to manually set fp16 in the network for the new apex. So here I set fp16=False. ################################################################################################################ ##这里是定义Es和G # 注意这里包含了两个步骤,Es编码+解码过程,既然解码(论文Figure 2的黄色梯形G)包含到这里了,下面Ea应该不会包含解码过程了 # 因为这里是一个类,如后续gen_a.encode()可以进行编码,gen_b.encode()可以进行解码 self.gen_a = AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen'], fp16=False) # auto-encoder for domain a self.gen_b = self.gen_a # auto-encoder for domain b ############################################################################################################################################ ############################################################################################################################################ ##这里是定义Ea # ID_stride,外观编码器池化层的stride if not 'ID_stride' in hyperparameters.keys(): hyperparameters['ID_stride'] = 2 # hyperparameters['ID_style']默认为'AB',论文中的Ea编码器 #这里是设置Ea,有三种模型可以选择 #PCB模型,ft_netAB为改造后的resnet50,ft_net为resnet50 if hyperparameters['ID_style'] == 'PCB': self.id_a = PCB(ID_class) elif hyperparameters['ID_style'] == 'AB': # 这是我们执行的模型,注意的是,id_a返回两个x(表示身份),获得f,具体介绍看函数内部 # 我们使用的是ft_netAB,是代码中Ea编码的过程,也就得到 ap code的过程,除了ap code还会得到两个分类结果 # 现在怀疑,该分类结果,可能就是行人重识别的结果 #ID_class表示有ID_class个不同ID的行人 self.id_a = ft_netAB(ID_class, stride=hyperparameters['ID_stride'], norm=hyperparameters['norm_id'], pool=hyperparameters['pool']) else: self.id_a = ft_net(ID_class, norm=hyperparameters['norm_id'], pool=hyperparameters['pool']) # return 2048 now # 这里进行的是浅拷贝,所以我认为他们的权重是一起的,可以理解为一个 self.id_b = self.id_a ############################################################################################################################################################ ############################################################################################################################################################ ##这里是定义D # 鉴别器,行人重识别,这里使用的是一个多尺寸的鉴别器,大概就是说,对图片进行几次缩放,并且对每次缩放都会预测,计算总的损失 # 经过网络3个元素,分别大小为[batch_size,1,64,32], [batch_size,1,32,16], [batch_size,1,16,8] self.dis_a = MsImageDis(3, hyperparameters['dis'], fp16=False) # discriminator for domain a self.dis_b = self.dis_a # discriminator for domain b ############################################################################################################################################################ ############################################################################################################################################################ # load teachers # 加载老师模型 # teacher:老师模型名称。对于DukeMTMC,您可以设置“best - duke” if hyperparameters['teacher'] != "": #teacher_name=best teacher_name = hyperparameters['teacher'] print(teacher_name) #有这个操作,我怀疑是可以加载多个教师模型 teacher_names = teacher_name.split(',') #构建老师模型 teacher_model = nn.ModuleList() teacher_count = 0 # 默认只有一个teacher_name='teacher_name',所以其加载的模型配置文件为项目根目录models/best/opts.yaml模型 for teacher_name in teacher_names: # 加载配置文件models/best/opts.yaml config_tmp = load_config(teacher_name) if 'stride' in config_tmp: #stride=1 stride = config_tmp['stride'] else: stride = 2 # 老师模型加载,老师模型为ft_net为resnet50 model_tmp = ft_net(ID_class, stride=stride) teacher_model_tmp = load_network(model_tmp, teacher_name) # 移除原本的全连接层 teacher_model_tmp.model.fc = nn.Sequential( ) # remove the original fc layer in ImageNet teacher_model_tmp = teacher_model_tmp.cuda() # summary(teacher_model_tmp, (3, 224, 224)) #使用浮点型 if self.fp16: teacher_model_tmp = amp.initialize(teacher_model_tmp, opt_level="O1") teacher_model.append(teacher_model_tmp.cuda().eval()) teacher_count += 1 self.teacher_model = teacher_model # 选择是否使用bn if hyperparameters['train_bn']: self.teacher_model = self.teacher_model.apply(train_bn) ############################################################################################################################################################ # 实例正则化 self.instancenorm = nn.InstanceNorm2d(512, affine=False) # RGB to one channel # 默认设置signal=gray,Es的输入为灰度图 if hyperparameters['single'] == 'edge': self.single = to_edge else: self.single = to_gray(False) # Random Erasing when training #earsing_p表示随机擦除的概率 if not 'erasing_p' in hyperparameters.keys(): self.erasing_p = 0 else: self.erasing_p = hyperparameters['erasing_p'] #随机擦除矩形区域的一些像素,应该类似于数据增强 self.single_re = RandomErasing(probability=self.erasing_p, mean=[0.0, 0.0, 0.0]) # 设置T_w为1,T_w为primary feature learning loss的权重系数 if not 'T_w' in hyperparameters.keys(): hyperparameters['T_w'] = 1 ################################################################################################ # Setup the optimizers # 设置优化器参数 beta1 = hyperparameters['beta1'] beta2 = hyperparameters['beta2'] dis_params = list( self.dis_a.parameters()) #+ list(self.dis_b.parameters()) gen_params = list( self.gen_a.parameters()) #+ list(self.gen_b.parameters()) #使用Adams优化器,用Adams训练Es,G,D self.dis_opt = torch.optim.Adam( [p for p in dis_params if p.requires_grad], lr=lr_d, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay']) self.gen_opt = torch.optim.Adam( [p for p in gen_params if p.requires_grad], lr=lr_g, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay']) # id params # 因为ID_style默认为AB,所以这里不执行 if hyperparameters['ID_style'] == 'PCB': ignored_params = ( list(map(id, self.id_a.classifier0.parameters())) + list(map(id, self.id_a.classifier1.parameters())) + list(map(id, self.id_a.classifier2.parameters())) + list(map(id, self.id_a.classifier3.parameters()))) base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters()) lr2 = hyperparameters['lr2'] #Ea 的优化器 self.id_opt = torch.optim.SGD( [{ 'params': base_params, 'lr': lr2 }, { 'params': self.id_a.classifier0.parameters(), 'lr': lr2 * 10 }, { 'params': self.id_a.classifier1.parameters(), 'lr': lr2 * 10 }, { 'params': self.id_a.classifier2.parameters(), 'lr': lr2 * 10 }, { 'params': self.id_a.classifier3.parameters(), 'lr': lr2 * 10 }], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True) # 这里是我们执行的代码 elif hyperparameters['ID_style'] == 'AB': # 忽略的参数,应该是适用于'PCB'或者其他的,但是不适用于'AB'的 ignored_params = ( list(map(id, self.id_a.classifier1.parameters())) + list(map(id, self.id_a.classifier2.parameters()))) # 获得基本的配置参数,如学习率 base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters()) lr2 = hyperparameters['lr2'] #对Ea使用SGD self.id_opt = torch.optim.SGD( [{ 'params': base_params, 'lr': lr2 }, { 'params': self.id_a.classifier1.parameters(), 'lr': lr2 * 10 }, { 'params': self.id_a.classifier2.parameters(), 'lr': lr2 * 10 }], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True) else: ignored_params = list(map(id, self.id_a.classifier.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters()) lr2 = hyperparameters['lr2'] self.id_opt = torch.optim.SGD( [{ 'params': base_params, 'lr': lr2 }, { 'params': self.id_a.classifier.parameters(), 'lr': lr2 * 10 }], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True) # 选择各个网络的优化 self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters) self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters) self.id_scheduler = get_scheduler(self.id_opt, hyperparameters) self.id_scheduler.gamma = hyperparameters['gamma2'] #ID Loss #交叉熵损失函数 self.id_criterion = nn.CrossEntropyLoss() # KL散度 self.criterion_teacher = nn.KLDivLoss(size_average=False) # Load VGG model if needed if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0: self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models') self.vgg.eval() for param in self.vgg.parameters(): param.requires_grad = False # save memory if self.fp16: # Name the FP16_Optimizer instance to replace the existing optimizer assert torch.backends.cudnn.enabled, "fp16 mode requires cudnn backend to be enabled." self.gen_a = self.gen_a.cuda() self.dis_a = self.dis_a.cuda() self.id_a = self.id_a.cuda() self.gen_b = self.gen_a self.dis_b = self.dis_a self.id_b = self.id_a self.gen_a, self.gen_opt = amp.initialize(self.gen_a, self.gen_opt, opt_level="O1") self.dis_a, self.dis_opt = amp.initialize(self.dis_a, self.dis_opt, opt_level="O1") self.id_a, self.id_opt = amp.initialize(self.id_a, self.id_opt, opt_level="O1")
def run(args): gpuId, epochs, weight_decay, batch_id, batch_image, lr_1, lr_2, erasing_p, sampling, exp_dir, trainset_name, cd_trainset_name, testset_names, rand_crop, head_1part_stride = \ args.gpuId, args.epochs, args.weight_decay, args.batch_id, args.batch_image, args.lr_1, args.lr_2, args.erasing_p, args.sampling, args.exp_dir, args.trainset_name, args.cd_trainset_name, args.testset_names, args.rand_crop, args.head_1part_stride DEVICE = torch.device("cuda:" + gpuId if torch.cuda.is_available() else "cpu") print(DEVICE) num_workers = 4 batch_test = 64 #32 train_list = [ transforms.Resize((400, 144)), transforms.RandomCrop((384, 128)) ] if rand_crop else [transforms.Resize((384, 128))] train_list += [ transforms.ToTensor(), ] if erasing_p > 0: train_list = train_list + [ RandomErasing(probability=erasing_p, mean=[0.0, 0.0, 0.0]) ] train_list += [ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ] train_transform = transforms.Compose(train_list) if trainset_name in ['market1501', 'cuhk03', 'duke']: root = get_dataset_root(trainset_name) train_dataset = Market1501( root + '/bounding_box_train', transform=train_transform, training=True, kpt_file=trainset_name + '-kpt.pkl' if args.pap else None, ps_dir=root + '_ps_label' if args.src_ps_lw > 0 else None) elif trainset_name in ['msmt17']: train_dataset = MSMT17(transform=train_transform, training=True, use_kpt=args.pap, use_ps=args.src_ps_lw > 0, split='train') else: raise ValueError('Invalid train set {}'.format(trainset_name)) train_loader = DataLoader(train_dataset, sampler=RandomIdSampler(train_dataset, batch_image=batch_image), batch_size=batch_id * batch_image, num_workers=num_workers, drop_last=True) if args.cd_ps_lw > 0: if cd_trainset_name in ['market1501', 'cuhk03', 'duke']: cd_train_dataset = Market1501(get_dataset_root(cd_trainset_name) + '/bounding_box_train', transform=train_transform, training=True, kpt_file=None, ps_dir=cd_trainset_name + '-ps') elif cd_trainset_name in ['msmt17']: cd_train_dataset = MSMT17(transform=train_transform, training=True, use_kpt=False, use_ps=True) else: raise ValueError( 'Invalid cd train set {}'.format(cd_trainset_name)) cd_train_loader = InfiniteNextBatch( DataLoader(cd_train_dataset, batch_size=args.cd_train_batch_size, num_workers=num_workers, drop_last=True)) test_transform = transforms.Compose([ transforms.Resize((384, 128)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_flip_transform = transforms.Compose([ transforms.Resize((384, 128)), functional.hflip, transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) def make_test_loader_M_C_D(root, name): query_dataset = Market1501(root + '/query', transform=test_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) query_flip_dataset = Market1501(root + '/query', transform=test_flip_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) query_loader = DataLoader(query_dataset, batch_size=batch_test, num_workers=num_workers) query_flip_loader = DataLoader(query_flip_dataset, batch_size=batch_test, num_workers=num_workers) test_dataset = Market1501(root + '/bounding_box_test', transform=test_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) test_flip_dataset = Market1501(root + '/bounding_box_test', transform=test_flip_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) test_loader = DataLoader(test_dataset, batch_size=batch_test, num_workers=num_workers) test_flip_loader = DataLoader(test_flip_dataset, batch_size=batch_test, num_workers=num_workers) return query_loader, query_flip_loader, test_loader, test_flip_loader def make_test_loader_MS_PR_PI(name): if name == 'msmt17': dclass = MSMT17 elif name == 'partial_reid': dclass = PartialREID elif name == 'partial_ilids': dclass = PartialiLIDs else: raise ValueError('Invalid dataset name {}'.format(name)) q_set = dclass(transform=test_transform, training=False, use_kpt=args.pap, use_ps=False, split='query') q_flip_set = dclass(transform=test_flip_transform, training=False, use_kpt=args.pap, use_ps=False, split='query') q_loader = DataLoader(q_set, batch_size=batch_test, num_workers=num_workers) q_flip_loader = DataLoader(q_flip_set, batch_size=batch_test, num_workers=num_workers) g_set = dclass(transform=test_transform, training=False, use_kpt=args.pap, use_ps=False, split='gallery') g_flip_set = dclass(transform=test_flip_transform, training=False, use_kpt=args.pap, use_ps=False, split='gallery') g_loader = DataLoader(g_set, batch_size=batch_test, num_workers=num_workers) g_flip_loader = DataLoader(g_flip_set, batch_size=batch_test, num_workers=num_workers) return q_loader, q_flip_loader, g_loader, g_flip_loader def make_test_loader(name): if name in ['market1501', 'cuhk03', 'duke']: return make_test_loader_M_C_D(get_dataset_root(name), name) elif name in ['msmt17', 'partial_reid', 'partial_ilids']: return make_test_loader_MS_PR_PI(name) test_loaders = [make_test_loader(name) for name in testset_names] mgn = MGN(len(train_dataset.unique_ids), args) if torch.cuda.device_count() > 1: mgn = nn.DataParallel(mgn) mgn = mgn.to(DEVICE) vanilla_cross_entropy_loss = nn.CrossEntropyLoss() cross_entropy_loss = nn.CrossEntropyLoss(reduce=False) triplet_semihard_loss = TripletSemihardLoss( margin=0.1, DEVICE=DEVICE, sampling=sampling, batch_id=batch_id, batch_image=batch_image) #batch_hard, .'curriculum' ps_loss = PSLoss() optimizer_start1 = optim.SGD(mgn.parameters(), lr=lr_1, momentum=0.9, weight_decay=weight_decay) optimizer_start2 = optim.SGD(mgn.parameters(), lr=lr_2, momentum=0.9, weight_decay=weight_decay) scheduler_1 = optim.lr_scheduler.MultiStepLR(optimizer_start1, [140, 180], gamma=0.1) scheduler_2 = optim.lr_scheduler.MultiStepLR( optimizer_start2, [140, 180], gamma=0.1) # best [140, 180] [120, 160] def get_model_input(inputs, target): dic = {'im': inputs.to(DEVICE)} if 'pap_mask_2p' in target: dic['pap_mask_2p'] = target['pap_mask_2p'].to(DEVICE) dic['pap_mask_3p'] = target['pap_mask_3p'].to(DEVICE) return dic def extract_loader_feat(loader, verbose=False): feat = [] vis = [] i = 0 for inputs, target in loader: if verbose: print(i) i += 1 with torch.no_grad(): output = mgn(get_model_input(inputs, target)) feat.append(output[1].detach().cpu().numpy()) if args.pap: vis_ = np.concatenate([ np.ones([len(output[1]), 3]), torch.stack(output[5 + 3 + 5:5 + 3 + 5 + 5], 1).detach().cpu().numpy() ], 1) vis.append(vis_) feat = np.concatenate(feat) vis = np.concatenate(vis) if args.pap else None return feat, vis def test(query_loader, query_flip_loader, test_loader, test_flip_loader, trainset_name, testset_name, epoch, verbose=False): cache_file = '{}/feat_cache-{}_to_{}.pkl'.format( exp_dir, trainset_name, testset_name) if args.use_feat_cache: assert os.path.exists( cache_file), "Feature cache file {} does not exist!".format( cache_file) query_2, q_vis, query_flip_2, q_vis, test_2, test_vis, test_flip_2, test_vis, q_ids, q_cams, g_ids, g_cams = load_pickle( cache_file) else: query_2, q_vis = extract_loader_feat(query_loader, verbose=verbose) query_flip_2, q_vis = extract_loader_feat(query_flip_loader, verbose=verbose) test_2, test_vis = extract_loader_feat(test_loader, verbose=verbose) test_flip_2, test_vis = extract_loader_feat(test_flip_loader, verbose=verbose) q_ids = query_loader.dataset.ids q_cams = query_loader.dataset.cameras g_ids = test_loader.dataset.ids g_cams = test_loader.dataset.cameras save_pickle([ query_2, q_vis, query_flip_2, q_vis, test_2, test_vis, test_flip_2, test_vis, q_ids, q_cams, g_ids, g_cams ], cache_file) if args.test_which_feat > 0: # TODO: implement for pap idx = args.test_which_feat query_2 = query_2[:, 256 * idx - 256:256 * idx] query_flip_2 = query_flip_2[:, 256 * idx - 256:256 * idx] test_2 = test_2[:, 256 * idx - 256:256 * idx] test_flip_2 = test_flip_2[:, 256 * idx - 256:256 * idx] query = normalize(query_2 + query_flip_2) test = normalize(test_2 + test_flip_2) if verbose: print('query.shape:', query.shape) print('test.shape:', test.shape) if args.pap: print('q_vis.shape:', q_vis.shape) print('test_vis.shape:', test_vis.shape) if args.pap: dist_1 = compute_dist_with_visibility(query, test, q_vis, test_vis, dist_type='euclidean', avg_by_vis_num=False) else: dist_1 = cdist(query, test) r_1 = cmc(dist_1, q_ids, g_ids, q_cams, g_cams, separate_camera_set=False, single_gallery_shot=False, first_match_break=True) m_ap_1 = mean_ap(dist_1, q_ids, g_ids, q_cams, g_cams) print('EPOCH [%d] %s -> %s: mAP=%f, r@1=%f, r@3=%f, r@5=%f, r@10=%f' % (epoch + 1, trainset_name, testset_name, m_ap_1, r_1[0], r_1[2], r_1[4], r_1[9])) if args.only_test: mgn.eval() if not args.use_feat_cache: if args.model_weight_file: model_weight_file = args.model_weight_file else: model_weight_file = '{}/model_weight.pth'.format(exp_dir) load_model_weight((mgn.module if hasattr(mgn, 'module') else mgn), model_weight_file) for name, test_loader in zip(testset_names, test_loaders): test(test_loader[0], test_loader[1], test_loader[2], test_loader[3], trainset_name, name, -1, verbose=False) exit() for epoch in range(epochs): mgn.train() scheduler_1.step() scheduler_2.step() running_loss = 0.0 running_loss_1 = 0.0 running_loss_2 = 0.0 if epoch < 20: optimizer_1 = optim.SGD(mgn.parameters(), lr=0.01 + 0.0045 * epoch, momentum=0.9, weight_decay=weight_decay) optimizer_2 = optim.SGD(mgn.parameters(), lr=0.001 + 0.00045 * epoch, momentum=0.9, weight_decay=weight_decay) else: optimizer_1 = optimizer_start1 optimizer_2 = optimizer_start2 for i, data in enumerate(train_loader): inputs, target = data inputs = inputs.to(DEVICE) for k, v in target.items(): target[k] = v.to(DEVICE) labels = target['id'] outputs = mgn(get_model_input(inputs, target)) optimizer_1.zero_grad() if args.pap: losses_1 = [ vanilla_cross_entropy_loss(output, labels) for output in outputs[5:5 + 3] ] + [(cross_entropy_loss(output, labels) * v).sum() / (v.sum() + 1e-12) for output, v in zip(outputs[5 + 3:5 + 3 + 5], outputs[5 + 3 + 5:5 + 3 + 5 + 5])] else: losses_1 = [ vanilla_cross_entropy_loss(output, labels) for output in outputs[5:5 + 8] ] loss_1 = sum(losses_1) / len(losses_1) psl = 0 if args.src_ps_lw > 0: psl = (ps_loss(outputs[-3], target['ps_label']) + ps_loss(outputs[-2], target['ps_label']) + ps_loss(outputs[-1], target['ps_label'])) / 3. (loss_1 + psl * args.src_ps_lw).backward() if args.cd_ps_lw > 0: cd_inputs, cd_targets = cd_train_loader.next_batch() cd_inputs = cd_inputs.to(DEVICE) for k, v in cd_targets.items(): cd_targets[k] = v.to(DEVICE) pap_old = args.pap args.pap = False outputs = mgn(get_model_input(cd_inputs, cd_targets)) args.pap = pap_old cd_psl = (ps_loss(outputs[-3], cd_targets['ps_label']) + ps_loss(outputs[-2], cd_targets['ps_label']) + ps_loss(outputs[-1], cd_targets['ps_label'])) / 3. (cd_psl * args.cd_ps_lw).backward() optimizer_1.step() outputs = mgn(get_model_input(inputs, target)) optimizer_2.zero_grad() losses_2 = [ triplet_semihard_loss(output, labels, epoch) for output in outputs[2:5] ] loss_2 = sum(losses_2) / len(losses_2) psl = 0 if args.src_ps_lw > 0: psl = (ps_loss(outputs[-3], target['ps_label']) + ps_loss(outputs[-2], target['ps_label']) + ps_loss(outputs[-1], target['ps_label'])) / 3. (loss_2 + psl * args.src_ps_lw).backward() if args.cd_ps_lw > 0: cd_inputs, cd_targets = cd_train_loader.next_batch() cd_inputs = cd_inputs.to(DEVICE) for k, v in cd_targets.items(): cd_targets[k] = v.to(DEVICE) pap_old = args.pap args.pap = False outputs = mgn(get_model_input(cd_inputs, cd_targets)) args.pap = pap_old cd_psl = (ps_loss(outputs[-3], cd_targets['ps_label']) + ps_loss(outputs[-2], cd_targets['ps_label']) + ps_loss(outputs[-1], cd_targets['ps_label'])) / 3. (cd_psl * args.cd_ps_lw).backward() optimizer_2.step() running_loss_1 += loss_1.item() running_loss_2 += loss_2.item() running_loss = running_loss + (loss_1.item() + loss_2.item()) / 2.0 print('%d/%d - %d/%d - loss: %f - ps_loss: %f - cd_ps_loss: %f' % (epoch + 1, epochs, i, len(train_loader), (loss_1.item() + loss_2.item()) / 2, psl.item() if isinstance(psl, torch.Tensor) else 0, cd_psl.item() if args.cd_ps_lw > 0 else 0)) print('epoch: %d/%d - loss1: %f' % (epoch + 1, epochs, running_loss_1 / len(train_loader))) print('epoch: %d/%d - loss2: %f' % (epoch + 1, epochs, running_loss_2 / len(train_loader))) # if (epoch + 1) % 50 == 0: # model_weight_file = '{}/model_weight.pth'.format(exp_dir) # save_model(mgn, model_weight_file) # mgn.eval() # for name, test_loader in zip(testset_names, test_loaders): # test(test_loader[0], test_loader[1], test_loader[2], test_loader[3], trainset_name, name, epoch) model_weight_file = '{}/model_weight.pth'.format(exp_dir) save_model(mgn, model_weight_file) mgn.eval() for name, test_loader in zip(testset_names, test_loaders): test(test_loader[0], test_loader[1], test_loader[2], test_loader[3], trainset_name, name, epoch)
# Load Data transform_train_list = [ transforms.Resize((384, 192), interpolation=transforms.InterpolationMode.BICUBIC), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] transform_val_list = [ transforms.Resize(size=(384, 192), interpolation=transforms.InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.erase_prob > 0: transform_train_list += [RandomErasing(probability=opt.erase_prob, mean=[0.0, 0.0, 0.0])] if opt.color_jitter: transform_train_list += [transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0)] print('dataset transformation: ', transform_train_list) data_transforms = { 'train': transforms.Compose(transform_train_list), 'val': transforms.Compose(transform_val_list), } image_datasets = { key: datasets.ImageFolder(os.path.join(data_dir, key), value) for key, value in data_transforms.items() }
random_height_crop = 224 transform_train_list = [ # transforms.RandomResizedCrop(size=128, scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC) transforms.Resize((288, 144), interpolation=3), transforms.RandomCrop((256, 128)), # transforms.Resize(256,interpolation=3), # transforms.RandomCrop(224,224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.erasing_p > 0: transform_train_list = transform_train_list + [ RandomErasing(opt.erasing_p) ] # print(transform_train_list) transform_val_list = [ transforms.Resize(size=(256, 128), interpolation=3), # Image.BICUBIC # transforms.Resize(256,interpolation=3), # transforms.RandomCrop(224,224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] data_transforms = { 'train': transforms.Compose(transform_train_list), 'val': transforms.Compose(transform_val_list),
transform_train_list = [ transforms.Resize((384, 128), interpolation=3), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] transform_val_list = [ transforms.Resize(size=(384, 128), interpolation=3), # Image.BICUBIC transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.erasing_p > 0: transform_train_list = transform_train_list + \ [RandomErasing(probability=opt.erasing_p, mean=[0.0, 0.0, 0.0])] if opt.color_jitter: transform_train_list = [ transforms.ColorJitter( brightness=0.1, contrast=0.1, saturation=0.1, hue=0) ] + transform_train_list print(transform_train_list) data_transforms = { 'train': transforms.Compose(transform_train_list), 'val': transforms.Compose(transform_val_list), } train_all = ''
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) # tensorboardX # writer = SummaryWriter(log_dir=osp.join(args.save_dir,'summary')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_img_dataset( root=args.root, name=args.dataset, split_id=args.split_id, cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split, ) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) if args.random_erasing: transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), RandomErasing(probability=args.probability, mean=[0.0, 0.0, 0.0]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False if args.loss == 'xent,htri': trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), sampler=RandomIdentitySampler(dataset.train, num_instances=args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) elif args.loss == 'xent': trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( ImageDataset(dataset.query, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( ImageDataset(dataset.gallery, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss=args.loss) print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters())/1000000.0)) criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) if args.stepsize > 0: if not args.warmup: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if args.resume: print("Loading checkpoint from '{}'".format(args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, use_gpu) return def adjust_lr(optimizer, ep): if ep < 20: lr = 1e-4 * (ep + 1) / 2 elif ep < 80: #lr = 1e-3 * len(args.gpu_devices) lr = 1e-3 elif ep < 180: #lr = 1e-4 * len(args.gpu_devices) lr = 1e-4 elif ep < 300: #lr = 1e-5 * len(args.gpu_devices) lr = 1e-5 elif ep < 320: #lr = 1e-5 * 0.1 ** ((ep - 320) / 80) * len(args.gpu_devices) lr = 1e-5 * 0.1 ** ((ep - 320) / 80) elif ep < 400: lr = 1e-6 elif ep < 480: #lr = 1e-4 * len(args.gpu_devices) lr = 1e-4 else: #lr = 1e-5 * len(args.gpu_devices) lr = 1e-5 for p in optimizer.param_groups: p['lr'] = lr length = len(trainloader) start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 #best_rerank1 = -np.inf #best_rerankepoch = 0 print("==> Start training") for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() if args.stepsize > 0: if args.warmup: adjust_lr(optimizer, epoch + 1) else: scheduler.step() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu=use_gpu, summary=None, length=length) train_time += round(time.time() - start_train_time) if (epoch+1) > args.start_eval and args.eval_step > 0 and (epoch+1) % args.eval_step == 0 or (epoch+1) == args.max_epoch: print("==> Test") rank1 = test(epoch, model, queryloader, galleryloader, use_gpu=True, summary=None) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 ####### Best Rerank #is_rerankbest = rerank1 > best_rerank1 #if is_rerankbest: # best_rerank1 = rerank1 # best_rerankepoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch+1) + '.pth.tar')) writer.close() print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(best_rank1, best_epoch)) #print("==> Best Rerank-1 {:.1%}, achieved at epoch {}".format(best_rerank1, best_rerankepoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time))
def __init__(self, hyperparameters, gpu_ids=[0]): super(DGNet_Trainer, self).__init__() lr_g = hyperparameters['lr_g'] #生成器学习率 lr_d = hyperparameters['lr_d'] #判别器学习率 ID_class = hyperparameters['ID_class'] if not 'apex' in hyperparameters.keys(): hyperparameters['apex'] = False self.fp16 = hyperparameters['apex'] # Initiate the networks # We do not need to manually set fp16 in the network for the new apex. So here I set fp16=False. # 构建Es编码+解码过程 gen_a.encode()可以进行编码,gen_b.encode()可以进行解码 self.gen_a = AdaINGen(hyperparameters['input_dim_a'], hyperparameters['gen'], fp16 = False) # auto-encoder for domain a self.gen_b = self.gen_a # auto-encoder for domain b # ID_stride,外观编码器池化层的stride if not 'ID_stride' in hyperparameters.keys(): hyperparameters['ID_stride'] = 2 # 构建外观编码器 if hyperparameters['ID_style']=='PCB': self.id_a = PCB(ID_class) elif hyperparameters['ID_style']=='AB': #使用的AB编码器 self.id_a = ft_netAB(ID_class, stride = hyperparameters['ID_stride'], norm=hyperparameters['norm_id'], pool=hyperparameters['pool']) else: self.id_a = ft_net(ID_class, norm=hyperparameters['norm_id'], pool=hyperparameters['pool']) # return 2048 now # 浅拷贝,两者等同 self.id_b = self.id_a # 鉴别器,使用的是一个多尺寸的鉴别器,即对图片进行几次缩放,并且对每次缩放都会预测,计算总的损失 self.dis_a = MsImageDis(3, hyperparameters['dis'], fp16 = False) # discriminator for domain a self.dis_b = self.dis_a # discriminator for domain b # load teachers 加载教师模型 if hyperparameters['teacher'] != "": teacher_name = hyperparameters['teacher'] print(teacher_name) # 构建教师模型 teacher_names = teacher_name.split(',') teacher_model = nn.ModuleList() teacher_count = 0 for teacher_name in teacher_names: config_tmp = load_config(teacher_name) if 'stride' in config_tmp: stride = config_tmp['stride'] else: stride = 2 # 网络搭建 model_tmp = ft_net(ID_class, stride = stride) teacher_model_tmp = load_network(model_tmp, teacher_name) teacher_model_tmp.model.fc = nn.Sequential() # remove the original fc layer in ImageNet teacher_model_tmp = teacher_model_tmp.cuda() if self.fp16: teacher_model_tmp = amp.initialize(teacher_model_tmp, opt_level="O1") teacher_model.append(teacher_model_tmp.cuda().eval()) teacher_count +=1 self.teacher_model = teacher_model # 选择是否使用bn if hyperparameters['train_bn']: self.teacher_model = self.teacher_model.apply(train_bn) # 实例正则化 self.instancenorm = nn.InstanceNorm2d(512, affine=False) # RGB to one channel if hyperparameters['single']=='edge': self.single = to_edge else: self.single = to_gray(False) # Random Erasing when training if not 'erasing_p' in hyperparameters.keys(): self.erasing_p = 0 else: self.erasing_p = hyperparameters['erasing_p'] # erasing_p表示随机擦除的概率 # 随机擦除矩形区域的一些像素,数据增强 self.single_re = RandomErasing(probability = self.erasing_p, mean=[0.0, 0.0, 0.0]) if not 'T_w' in hyperparameters.keys(): hyperparameters['T_w'] = 1 # Setup the optimizers 设置优化器的参数 beta1 = hyperparameters['beta1'] beta2 = hyperparameters['beta2'] dis_params = list(self.dis_a.parameters()) #+ list(self.dis_b.parameters()) gen_params = list(self.gen_a.parameters()) #+ list(self.gen_b.parameters()) # 使用Adam优化器 self.dis_opt = torch.optim.Adam([p for p in dis_params if p.requires_grad], lr=lr_d, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay']) self.gen_opt = torch.optim.Adam([p for p in gen_params if p.requires_grad], lr=lr_g, betas=(beta1, beta2), weight_decay=hyperparameters['weight_decay']) # id params if hyperparameters['ID_style']=='PCB': ignored_params = (list(map(id, self.id_a.classifier0.parameters() )) +list(map(id, self.id_a.classifier1.parameters() )) +list(map(id, self.id_a.classifier2.parameters() )) +list(map(id, self.id_a.classifier3.parameters() )) ) base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters()) lr2 = hyperparameters['lr2'] self.id_opt = torch.optim.SGD([ {'params': base_params, 'lr': lr2}, {'params': self.id_a.classifier0.parameters(), 'lr': lr2*10}, {'params': self.id_a.classifier1.parameters(), 'lr': lr2*10}, {'params': self.id_a.classifier2.parameters(), 'lr': lr2*10}, {'params': self.id_a.classifier3.parameters(), 'lr': lr2*10} ], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True) elif hyperparameters['ID_style']=='AB': ignored_params = (list(map(id, self.id_a.classifier1.parameters())) + list(map(id, self.id_a.classifier2.parameters()))) # 获得基本的配置参数,如学习率 base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters()) lr2 = hyperparameters['lr2'] self.id_opt = torch.optim.SGD([ {'params': base_params, 'lr': lr2}, {'params': self.id_a.classifier1.parameters(), 'lr': lr2*10}, {'params': self.id_a.classifier2.parameters(), 'lr': lr2*10} ], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True) else: ignored_params = list(map(id, self.id_a.classifier.parameters() )) base_params = filter(lambda p: id(p) not in ignored_params, self.id_a.parameters()) lr2 = hyperparameters['lr2'] self.id_opt = torch.optim.SGD([ {'params': base_params, 'lr': lr2}, {'params': self.id_a.classifier.parameters(), 'lr': lr2*10} ], weight_decay=hyperparameters['weight_decay'], momentum=0.9, nesterov=True) # 选择各个网络优化的策略 self.dis_scheduler = get_scheduler(self.dis_opt, hyperparameters) self.gen_scheduler = get_scheduler(self.gen_opt, hyperparameters) self.id_scheduler = get_scheduler(self.id_opt, hyperparameters) self.id_scheduler.gamma = hyperparameters['gamma2'] #ID Loss self.id_criterion = nn.CrossEntropyLoss() self.criterion_teacher = nn.KLDivLoss(size_average=False) # Load VGG model if needed if 'vgg_w' in hyperparameters.keys() and hyperparameters['vgg_w'] > 0: self.vgg = load_vgg16(hyperparameters['vgg_model_path'] + '/models') self.vgg.eval() for param in self.vgg.parameters(): param.requires_grad = False # save memory if self.fp16: # Name the FP16_Optimizer instance to replace the existing optimizer assert torch.backends.cudnn.enabled, "fp16 mode requires cudnn backend to be enabled." self.gen_a = self.gen_a.cuda() self.dis_a = self.dis_a.cuda() self.id_a = self.id_a.cuda() self.gen_b = self.gen_a self.dis_b = self.dis_a self.id_b = self.id_a self.gen_a, self.gen_opt = amp.initialize(self.gen_a, self.gen_opt, opt_level="O1") self.dis_a, self.dis_opt = amp.initialize(self.dis_a, self.dis_opt, opt_level="O1") self.id_a, self.id_opt = amp.initialize(self.id_a, self.id_opt, opt_level="O1")
if opt.PCB: transform_train_list = [ transforms.Resize((384, 192), interpolation=3), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] transform_val_list = [ transforms.Resize(size=(384, 192), interpolation=3), #Image.BICUBIC transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.erasing_p > 0: transform_train_list = transform_train_list + [ RandomErasing(probability=opt.erasing_p, mean=[0.0, 0.0, 0.0]) ] if opt.color_jitter: transform_train_list = [ transforms.ColorJitter( brightness=0.1, contrast=0.1, saturation=0.1, hue=0) ] + transform_train_list print('transform_train list =', transform_train_list) data_transforms = { 'train': transforms.Compose(transform_train_list), 'val': transforms.Compose(transform_val_list), } train_all = ''
# useful for sampler self.lb_img_dict = dict() self.lb_ids_uniq = set(self.lb_ids) lb_array = np.array(self.lb_ids) for lb in self.lb_ids_uniq: idx = np.where(lb_array == lb)[0] self.lb_img_dict.update({lb: idx}) def __len__(self): return len(self.imgs) def __getitem__(self, idx): img = Image.open(self.imgs[idx]) img = self.trans(img) return img, self.lb_ids[idx], self.lb_cams[idx], self.imgs[idx] if __name__ == "__main__": img_dir = "/home/CORP/ryann.bai/dataset/VehicleID/image/" img_list = "/home/CORP/ryann.bai/dataset/VehicleID/train_test_split_v1/train_list_start0_jpg.txt" ds = VehicleID(img_dir, img_list, is_train=True) im, _, _ = ds[1] print(im.shape) print(im.max()) print(im.min()) ran_er = RandomErasing() im = ran_er(im) cv2.imshow('erased', im) cv2.waitKey(0)
import os import numpy as np from sklearn.model_selection import train_test_split from runstats import Statistics import torch from torch.utils.data import Dataset from torchvision import transforms from albumentations import Compose, ShiftScaleRotate, GridDistortion from albumentations.pytorch import ToTensor from random_erasing import RandomErasing random_erasing = RandomErasing() albumentations_transform = Compose([ ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=0.5), GridDistortion(), ToTensor() ]) def get_train_val_data(inp_txt): with open(inp_txt, 'r') as f: lines = f.readlines() lines = [line.rstrip() for line in lines] fnames = [line.split('\t')[0] for line in lines] labels = [line.split('\t')[-1] for line in lines] ref_labels = list(sorted(set(labels)))
def train(opt): version = torch.__version__ fp16 = opt.fp16 data_dir = opt.data_dir name = opt.name str_ids = opt.gpu_ids.split(',') gpu_ids = [] for str_id in str_ids: gid = int(str_id) if gid >= 0: gpu_ids.append(gid) # set gpu ids if len(gpu_ids) > 0: torch.cuda.set_device(gpu_ids[0]) cudnn.benchmark = True ###################################################################### # Load Data # --------- # transform_train_list = [ # transforms.RandomResizedCrop(size=128, scale=(0.75,1.0), ratio=(0.75,1.3333), interpolation=3), #Image.BICUBIC) transforms.Resize((256, 128), interpolation=3), transforms.Pad(10), transforms.RandomCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] transform_val_list = [ transforms.Resize(size=(256, 128), interpolation=3), # Image.BICUBIC transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.PCB: transform_train_list = [ transforms.Resize((384, 192), interpolation=3), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] transform_val_list = [ transforms.Resize(size=(384, 192), interpolation=3), # Image.BICUBIC transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.erasing_p > 0: transform_train_list = transform_train_list + [ RandomErasing(probability=opt.erasing_p, mean=[0.0, 0.0, 0.0]) ] if opt.color_jitter: transform_train_list = [ transforms.ColorJitter( brightness=0.1, contrast=0.1, saturation=0.1, hue=0) ] + transform_train_list # print(transform_train_list) data_transforms = { 'train': transforms.Compose(transform_train_list), 'val': transforms.Compose(transform_val_list), } train_all = '' if opt.train_all: train_all = '_all' image_datasets = {} image_datasets['train'] = datasets.ImageFolder( os.path.join(data_dir, 'train' + train_all), data_transforms['train']) image_datasets['val'] = datasets.ImageFolder(os.path.join(data_dir, 'val'), data_transforms['val']) dataloaders = { x: torch.utils.data.DataLoader(image_datasets[x], batch_size=opt.batchsize, shuffle=True, num_workers=8, pin_memory=True) # 8 workers may work faster for x in ['train', 'val'] } dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} class_names = image_datasets['train'].classes use_gpu = torch.cuda.is_available() #since = time.time() #inputs, classes = next(iter(dataloaders['train'])) #print('time used for loading data: %ds' %(time.time() - since)) ###################################################################### # Training the model # ------------------ # # Now, let's write a general function to train a model. Here, we will # illustrate: # # - Scheduling the learning rate # - Saving the best model # # In the following, parameter ``scheduler`` is an LR scheduler object from # ``torch.optim.lr_scheduler``. y_loss = {} # loss history y_loss['train'] = [] y_loss['val'] = [] y_err = {} y_err['train'] = [] y_err['val'] = [] def train_model(model, criterion, optimizer, scheduler, num_epochs=25): since = time.time() results = [] for epoch in range(num_epochs): print('Epoch {}/{}'.format(epoch, num_epochs - 1)) # Each epoch has a training and validation phase for phase in ['train', 'val']: if phase == 'train': scheduler.step() model.train(True) # Set model to training mode else: model.train(False) # Set model to evaluate mode running_loss = 0.0 running_corrects = 0.0 # Iterate over data. pbar = tqdm(dataloaders[phase]) for inputs, labels in pbar: # get the inputs now_batch_size, c, h, w = inputs.shape if now_batch_size < opt.batchsize: # skip the last batch continue # print(inputs.shape) # wrap them in Variable if use_gpu: inputs = Variable(inputs.cuda().detach()) labels = Variable(labels.cuda().detach()) else: inputs, labels = Variable(inputs), Variable(labels) # if we use low precision, input also need to be fp16 # if fp16: # inputs = inputs.half() # zero the parameter gradients optimizer.zero_grad() # forward if phase == 'val': with torch.no_grad(): outputs = model(inputs) else: outputs = model(inputs) if not opt.PCB: _, preds = torch.max(outputs.data, 1) loss = criterion(outputs, labels) else: part = {} sm = nn.Softmax(dim=1) num_part = 6 for i in range(num_part): part[i] = outputs[i] score = sm(part[0]) + sm(part[1]) + sm(part[2]) + sm( part[3]) + sm(part[4]) + sm(part[5]) _, preds = torch.max(score.data, 1) loss = criterion(part[0], labels) for i in range(num_part - 1): loss += criterion(part[i + 1], labels) # backward + optimize only if in training phase if phase == 'train': if fp16: # we use optimier to backward loss with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() optimizer.step() # statistics if int(version[0]) > 0 or int( version[2] ) > 3: # for the new version like 0.4.0, 0.5.0 and 1.0.0 running_loss += loss.item() * now_batch_size else: # for the old version like 0.3.0 and 0.3.1 running_loss += loss.data[0] * now_batch_size running_corrects += float(torch.sum(preds == labels.data)) pbar.set_description( desc='loss: {:.4f}'.format(loss.item())) epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects / dataset_sizes[phase] print('\r\n{} Loss: {:.4f} Acc: {:.4f}'.format( phase, epoch_loss, epoch_acc)) logging.info('epoch: {}, {} Loss: {:.4f} Acc: {:.4f}'.format( epoch, phase, epoch_loss, epoch_acc)) y_loss[phase].append(epoch_loss) y_err[phase].append(1.0 - epoch_acc) # deep copy the model if phase == 'val': results.append({ 'epoch': epoch, 'trainLoss': y_loss['train'][-1], 'trainError': y_err['train'][-1], 'valLoss': y_loss['val'][-1], 'valError': y_err['val'][-1] }) last_model_wts = model.state_dict() if epoch % 10 == 9: save_network(model, epoch) draw_curve(epoch) write_to_csv(results) time_elapsed = time.time() - since print('\r\nTraining complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) print() time_elapsed = time.time() - since print('\r\nTraining complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) # print('Best val Acc: {:4f}'.format(best_acc)) # load best model weights model.load_state_dict(last_model_wts) save_network(model, 'last') return model ###################################################################### # Draw Curve # --------------------------- x_epoch = [] fig = plt.figure() ax0 = fig.add_subplot(121, title="loss") ax1 = fig.add_subplot(122, title="top1err") def draw_curve(current_epoch): x_epoch.append(current_epoch) ax0.plot(x_epoch, y_loss['train'], 'bo-', label='train') ax0.plot(x_epoch, y_loss['val'], 'ro-', label='val') ax1.plot(x_epoch, y_err['train'], 'bo-', label='train') ax1.plot(x_epoch, y_err['val'], 'ro-', label='val') if current_epoch == 0: ax0.legend() ax1.legend() fig.savefig(os.path.join('./model', name, 'train.jpg')) def write_to_csv(results): path = os.path.join('./model', name, 'result.csv') with open(path, 'w', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=list(results[0].keys())) writer.writeheader() writer.writerows(results) ###################################################################### # Save model # --------------------------- def save_network(network, epoch_label): save_filename = 'net_%s.pth' % epoch_label rpth = os.path.join('./model', name, 'Model Files') if not os.path.exists(rpth): os.makedirs(rpth) save_path = os.path.join(rpth, save_filename) torch.save(network.cpu().state_dict(), save_path) if torch.cuda.is_available(): network.cuda(gpu_ids[0]) ###################################################################### # Finetuning the convnet # ---------------------- # # Load a pretrainied model and reset final fully connected layer. # if opt.use_dense: model = ft_net_dense(len(class_names), opt.droprate) else: model = ft_net(len(class_names), opt.droprate, opt.stride) if opt.PCB: model = PCB(len(class_names)) opt.nclasses = len(class_names) print(model) print('model loaded') if not opt.PCB: ignored_params = list(map(id, model.model.fc.parameters())) + list( map(id, model.classifier.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer_ft = optim.SGD([{ 'params': base_params, 'lr': 0.1 * opt.lr }, { 'params': model.model.fc.parameters(), 'lr': opt.lr }, { 'params': model.classifier.parameters(), 'lr': opt.lr }], weight_decay=5e-4, momentum=0.9, nesterov=True) else: ignored_params = list(map(id, model.model.fc.parameters())) ignored_params += ( list(map(id, model.classifier0.parameters())) + list(map(id, model.classifier1.parameters())) + list(map(id, model.classifier2.parameters())) + list(map(id, model.classifier3.parameters())) + list(map(id, model.classifier4.parameters())) + list(map(id, model.classifier5.parameters())) # +list(map(id, model.classifier6.parameters() )) # +list(map(id, model.classifier7.parameters() )) ) base_params = filter(lambda p: id(p) not in ignored_params, model.parameters()) optimizer_ft = optim.SGD( [ { 'params': base_params, 'lr': 0.1 * opt.lr }, { 'params': model.model.fc.parameters(), 'lr': opt.lr }, { 'params': model.classifier0.parameters(), 'lr': opt.lr }, { 'params': model.classifier1.parameters(), 'lr': opt.lr }, { 'params': model.classifier2.parameters(), 'lr': opt.lr }, { 'params': model.classifier3.parameters(), 'lr': opt.lr }, { 'params': model.classifier4.parameters(), 'lr': opt.lr }, { 'params': model.classifier5.parameters(), 'lr': opt.lr }, # {'params': model.classifier6.parameters(), 'lr': 0.01}, # {'params': model.classifier7.parameters(), 'lr': 0.01} ], weight_decay=5e-4, momentum=0.9, nesterov=True) # Decay LR by a factor of 0.1 every 40 epochs exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=40, gamma=0.1) ###################################################################### # Train and evaluate # ^^^^^^^^^^^^^^^^^^ # # It should take around 1-2 hours on GPU. # dir_name = os.path.join('./model', name) if not os.path.isdir(dir_name): os.mkdir(dir_name) # record every run copyfile('./train.py', dir_name + '/train.py') copyfile('./model.py', dir_name + '/model.py') # save opts with open('%s/opts.yaml' % dir_name, 'w') as fp: yaml.dump(vars(opt), fp, default_flow_style=False) # model to gpu model = model.cuda() if fp16: # model = network_to_half(model) # optimizer_ft = FP16_Optimizer(optimizer_ft, static_loss_scale = 128.0) model, optimizer_ft = amp.initialize(model, optimizer_ft, opt_level="O1") criterion = losses.DualLoss() model = train_model(model, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=60) # # if __name__ == "__main__": # train(opt)
train_preprocessing = [ transforms.Resize((288, 144), interpolation=3), transforms.RandomCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] val_preprocessing = [ transforms.Resize((256, 128), interpolation=3), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if ag.erasing_p > 0: train_preprocessing = train_preprocessing + [RandomErasing(probability=ag.erasing_p, mean=[0.0, 0.0, 0.0])] if ag.color_jitter: train_preprocessing = [transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0)]\ + train_preprocessing # Compose all the transformations data_preprocessing = { 'train': transforms.Compose(train_preprocessing), 'val': transforms.Compose(val_preprocessing), } # Use all the training samples or not train_all = '' if ag.train_all: train_all = '_all'
if opt.PCB: transform_train_list = [ transforms.Resize((384, 192), interpolation=3), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] transform_val_list = [ transforms.Resize(size=(384, 192), interpolation=3), # Image.BICUBIC transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if opt.erasing_p > 0: transform_train_list = transform_train_list + [RandomErasing(probability=opt.erasing_p, mean=[0.0, 0.0, 0.0])] if opt.color_jitter: transform_train_list = [transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0)] + transform_train_list print(transform_train_list) data_transforms = { 'train': transforms.Compose(transform_train_list), 'val': transforms.Compose(transform_val_list), } train_all = '' if opt.train_all: train_all = '_all'
args = parser.parse_args() print('Loading data') transform_list = [ transforms.Resize((256, 128), interpolation=3), transforms.Pad(10), transforms.RandomCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] if args.erasing_p > 0: transform_list.append( RandomErasing(probability=args.erasing_p, mean=[0.0, 0.0, 0.0])) if args.color_jitter: transform_list.append( transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0)) # Load dataset transform_list = transforms.Compose(transform_list) dset_train = MultiViewDataSet(args.data, transform=transform_list) train_loader = DataLoader(dset_train, batch_size=args.batch_size, shuffle=True, num_workers=8, pin_memory=True)