def checkpoint(self, ts=None, best=False): if ts is None: ts = self.ts filename = ts + 'checkpoint.pth.tar' filename_best = ts + 'best_checkpoint.pth.tar' state = { 'state_dicts': [model.state_dict() for model in self.model_list], 'epoch': self.epoch + 1 if not self.args.evaluate else None, 'optimizers': [optimizer.state_dict() for optimizer in self.optimizer_list], 'losses': self.epoch_losses if not self.args.evaluate else None, 'performance_evolution': self.test_acc_logger.get_data(), 'args': self.loaded_args if self.args.evaluate else self.args, 'ts': self.ts, 'result_dict': self.result_dict } save_checkpoint(state, osp.join(self.args.save_experiment, filename)) if best: save_checkpoint(state, osp.join(self.args.save_experiment, filename_best)) print("Saved model checkpoint at " + filename) if best: print("Also saved as best checkpoint. ")
def run(self): start_time = time.time() train_time = 0 best_epoch = 0 best_acc = self._eval_epoch(epoch=0) # === Train for epoch in range(1, self.args.max_epoch + 1): self.lr = self.optimizer.param_groups[0]['lr'] self.writer.add_scalar('Learning rate', self.lr, global_step=epoch) epoch_time_start = time.time() self._train_epoch(epoch) train_time += round(time.time() - epoch_time_start) # === Evaluation if epoch == 1 or epoch == self.args.max_epoch or epoch % self.args.eval_per_epoch == 0: acc = self._eval_epoch(epoch) is_best = acc > best_acc if is_best: best_acc = acc best_epoch = epoch if is_best or epoch % self.args.save_per_epoch == 0: state = { 'encoder_state_dict': self.encoder.state_dict(), 'decoder_state_dict': self.decoder.state_dict(), 'acc': acc, 'epoch': epoch, } file_path = os.path.join( self.args.save_dir, self.dataset, 'pretrain', self.args.encoder, 'ckp_ep' + str(epoch) + '.ckp') save_checkpoint(state, is_best, file_path) print('==> Test best accuracy {:.2%}, achieved at epoch {}'.format( best_acc, best_epoch)) torch.cuda.empty_cache() elapsed = str(datetime.timedelta(seconds=round(time.time() - start_time))) 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)) print("==========\nArgs:{}\n==========".format(self.args))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) train_dataset = data_manager.init_dataset(root=args.root, name='json', phase='train') valid_dataset = data_manager.init_dataset(root=args.root, name='json', phase='valid') test_dataset = data_manager.init_dataset(root=args.root, name='json', phase='test') test_mask = test_dataset.mask pin_memory = True if use_gpu else False trainloader = DataLoader(JsonDataset(train_dataset), num_workers=4, batch_size=args.train_batch, pin_memory=pin_memory, drop_last=True) validloader = DataLoader(JsonDataset(valid_dataset), num_workers=4, batch_size=args.test_batch, pin_memory=pin_memory, drop_last=True) testloader = DataLoader(JsonDataset(test_dataset), num_workers=4, batch_size=args.test_batch, pin_memory=pin_memory, drop_last=True) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, use_gpu=use_gpu) print("Model size: {:.3f} M".format(count_num_param(model))) model.init_weights() criterion_label = LabelLoss() optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if args.load_weights: # load pretrained weights but ignore layers that don't match in size print("Loading pretrained weights from '{}'".format(args.load_weights)) checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) if args.resume: checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}".format(start_epoch)) # if use_gpu: # str_ids = args.gpu_devices.split(',') # gpu_ids = [] # for str_id in str_ids: # id = int(str_id) # if id >= 0: # gpu_ids.append(id) # model = nn.DataParallel(model, gpu_ids) device = torch.device( 'cuda' if use_gpu and torch.cuda.is_available() else 'cpu') model.to(device) if args.evaluate: print("Evaluate only") start_evaluate_time = time.time() test_thetas = evaluate(model, testloader, use_gpu, args.test_batch, test_mask) # test_thetas = evaluate(model, validloader, use_gpu, args.test_batch, test_mask) evaluate_time = time.time() - start_evaluate_time print('Evaluate: {} secs'.format(evaluate_time)) with open("auto_sample.csv", "r") as csvfiler: with open("test_thetas.csv", "w") as csvfilew: reader = csv.reader(csvfiler) for item in reader: if reader.line_num == 1: writer = csv.writer(csvfilew) writer.writerow(['test_id', 'result']) continue writer = csv.writer(csvfilew) writer.writerow( [item[0], str(test_thetas[reader.line_num - 2])]) # writer.writerows(map(lambda x: [x], test_thetas)) return start_time = time.time() train_time = 0 best_label_loss = np.inf best_epoch = 0 # # print("==> Test") # label_loss = test(model, validloader, criterion_label, use_gpu, args.test_batch) # print("test label loss RMES() is {}".format(label_loss)) # # print("==> Start training") for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_label, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() # save model every epoch if (epoch + 1) % args.save_step == 0: print("==> Now save epoch {} \'s model".format(epoch + 1)) # if use_gpu: # state_dict = model.state_dict() # module. # else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'epoch': epoch }, False, osp.join(args.save_dir, 'checkpoint_latest.pth')) # test model every eval_step if (epoch + 1) > args.start_eval and args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or \ (epoch + 1) == args.max_epoch: print("==> Test") label_loss = test(model, validloader, criterion_label, use_gpu, args.test_batch) is_best = label_loss < best_label_loss if is_best: best_label_loss = label_loss best_epoch = epoch + 1 # if use_gpu: # state_dict = model.state_dict() # else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'label_loss': label_loss, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth')) # .pth.tar print("==> Best Label Loss {:.3}, achieved at epoch {}".format( best_label_loss, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False 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, use_lmdb=args.use_lmdb, ) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train, use_lmdb=args.use_lmdb, lmdb_path=dataset.train_lmdb_path), 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, use_lmdb=args.use_lmdb, lmdb_path=dataset.query_lmdb_path), 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, use_lmdb=args.use_lmdb, lmdb_path=dataset.gallery_lmdb_path), 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))) criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if args.fixbase_epoch > 0: if hasattr(model, 'classifier') and isinstance(model.classifier, nn.Module): optimizer_tmp = init_optim(args.optim, model.classifier.parameters(), args.fixbase_lr, args.weight_decay) else: print("Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0") args.fixbase_epoch = 0 if args.load_weights: # load pretrained weights but ignore layers that don't match in size print("Loading pretrained weights from '{}'".format(args.load_weights)) checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()} model_dict.update(pretrain_dict) model.load_state_dict(model_dict) if args.resume: checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, use_gpu) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") if args.fixbase_epoch > 0: print("Train classifier for {} epochs while keeping base network frozen".format(args.fixbase_epoch)) for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer_tmp, trainloader, use_gpu, freeze_bn=True) train_time += round(time.time() - start_train_time) del optimizer_tmp print("Now open all layers for training") for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch+1) > args.start_eval and args.eval_step > 0 and (epoch+1) % args.eval_step == 0 or (epoch+1) == args.max_epoch: print("==> Test") rank1 = test(model, queryloader, galleryloader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch+1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_vidreid_dataset(root=args.root, name=args.dataset) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False # decompose tracklets into images for image-based training new_train = [] for img_paths, pid, camid in dataset.train: for img_path in img_paths: new_train.append((img_path, pid, camid)) trainloader = DataLoader( ImageDataset(new_train, transform=transform_train), sampler=RandomIdentitySampler(new_train, num_instances=args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( VideoDataset(dataset.query, seq_len=args.seq_len, sample='evenly', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( VideoDataset(dataset.gallery, seq_len=args.seq_len, sample='evenly', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(model))) 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) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=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, args.pool, use_gpu) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, queryloader, galleryloader, args.pool, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset( root=args.root, name=args.dataset, split_id=args.split_id, cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split, use_lmdb=args.use_lmdb, ) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train, use_lmdb=args.use_lmdb, lmdb_path=dataset.train_lmdb_path), 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, ) queryloader = DataLoader( ImageDataset(dataset.query, transform=transform_test, use_lmdb=args.use_lmdb, lmdb_path=dataset.train_lmdb_path), 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, use_lmdb=args.use_lmdb, lmdb_path=dataset.train_lmdb_path), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(model))) 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) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.load_weights: # load pretrained weights but ignore layers that don't match in size if check_isfile(args.load_weights): checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format( args.load_weights)) if args.resume: if check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format( args.start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.vis_ranked_res: visualize_ranked_results( distmat, dataset, save_dir=osp.join(args.save_dir, 'ranked_results'), topk=20, ) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, queryloader, galleryloader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset( root=args.root, name=args.dataset, split_id=args.split_id, cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split, use_lmdb=args.use_lmdb, ) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), # T.Resize((args.height, args.width)), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainset = ImageDataset(dataset.train, transform=transform_train, use_lmdb=args.use_lmdb, lmdb_path=dataset.train_lmdb_path) trainloader = DataLoader( trainset, batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryset = ImageDataset(dataset.query, transform=transform_test, use_lmdb=args.use_lmdb, lmdb_path=dataset.query_lmdb_path) queryloader = DataLoader( queryset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryset = ImageDataset(dataset.gallery, transform=transform_test, use_lmdb=args.use_lmdb, lmdb_path=dataset.gallery_lmdb_path) galleryloader = DataLoader( galleryset, 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))) # summary(model, (3, 160, 64)) criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if args.fixbase_epoch > 0: if hasattr(model, 'classifier') and isinstance(model.classifier, nn.Module): optimizer_tmp = init_optim(args.optim, model.classifier.parameters(), args.fixbase_lr, args.weight_decay) else: print( "Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0" ) args.fixbase_epoch = 0 if args.load_weights: # load pretrained weights but ignore layers that don't match in size print("Loading pretrained weights from '{}'".format(args.load_weights)) if torch.cuda.is_available(): checkpoint = torch.load(args.load_weights) else: checkpoint = torch.load(args.load_weights, map_location='cpu') pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) if args.resume: checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, use_gpu) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") if args.fixbase_epoch > 0: print( "Train classifier for {} epochs while keeping base network frozen". format(args.fixbase_epoch)) for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer_tmp, trainloader, use_gpu, freeze_bn=True) train_time += round(time.time() - start_train_time) del optimizer_tmp print("Now open all layers for training") for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, queryloader, galleryloader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) ''' if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'epoch': epoch, }, True, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch+1) + '.pth.tar')) ''' print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'), mode='a') else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'), mode='a') 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(name=args.dataset, dataset_dir=args.root, fore_dir=args.fore_dir) transform_train = ST.Compose([ ST.Scale((args.height, args.width), interpolation=3), ST.RandomHorizontalFlip(), ST.ToTensor(), ST.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ST.RandomErasing(0.5) ]) transform_test = ST.Compose([ ST.Scale((args.height, args.width), interpolation=3), ST.ToTensor(), ST.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset_hardSplit_seg(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, ) 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) print(model) criterion_xent = CrossEntropyLabelSmooth(use_gpu=use_gpu) criterion_htri = TripletLoss() criterion_mask = MaskLoss() criterion_split = HardSplitLoss() criterion_cluster = ClusterLoss() 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.resume: if check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format( args.start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, use_gpu) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, criterion_mask, criterion_split, criterion_cluster, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and ( epoch + 1) % args.eval_step == 0 or epoch == 0: 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)) print("==========\nArgs:{}\n==========".format(args))
def _train_session(self, session, use_centroid=False): # assert session in range(0, 9) print('Training session {}'.format(session)) self.cur_session = session if session > 0: self.memory.del_noisy_slots() memory_vals = self.memory.m_vals.cpu().numpy() memory_vals_counter = Counter(memory_vals) print('memory val:', len(memory_vals_counter), memory_vals_counter.most_common()) # === Data print('Preparing data {} with session {}...'.format(self.dataset, session)) train_loader, eval_loader = self.get_dataloader(session) if session == 0: # train_loader, eval_loader = self.data_manager.get_dataloaders() max_epoch = self.args.max_epoch_sess0 m_replays = None else: # if self.dataset == 'miniImageNet': # train_loader = self.data_manager.get_dataloaders(session=session, is_fewshot=True) # eval_loader = self.data_manager.get_dataloaders(session=session, is_fewshot=False) # elif self.dataset == 'cifar100' or 'cub-200' or 'casia-face': # train_loader, eval_loader = self.data_manager.get_dataloaders( # session=session, is_fewshot=True) # else: # raise NotImplementedError # Memory replay data m_keys, m_targets = self._get_nonempty_memory_slots() # m_keys = m_keys.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, 3, 3) # [nonempty, key_dim, 3, 3] m_replays = (m_keys, m_targets) max_epoch = 20 # Todo print('Num of batches of train loader:', len(train_loader)) self.eval_dataloaders.append(eval_loader) start_time = time.time() train_time = 0 best_epoch = 0 # best_acc = self._eval_epoch(epoch=0, eval_loader=eval_loader) best_acc = -1 best_state = None # === Train for epoch in range(1, max_epoch + 1): if session == 0: if epoch == 1: adjust_lr(self.optimizer_all, 1e-3) if epoch > 10: adjust_lr(self.optimizer_all, 1e-4) else: if epoch == 1: adjust_lr(self.optimizer_all, 1e-4) if epoch > 15: adjust_lr(self.optimizer_all, 1e-5) cur_lr = self.optimizer_all.param_groups[0]['lr'] self.writer.add_scalar('Learning rate', cur_lr, global_step=epoch) epoch_time_start = time.time() self._train_epoch(epoch, cur_lr, train_loader, m_replays) train_time += round(time.time() - epoch_time_start) # === Eval on current session's dataloader only if epoch == self.args.max_epoch_sess0 or epoch % self.args.eval_per_epoch == 0: if use_centroid: self.memory.upd_centroids() acc = self._eval_epoch(epoch, eval_loader, use_centroid=use_centroid) # === Save checkpoint is_best = acc > best_acc if is_best or epoch == 2 or epoch % self.args.save_per_epoch == 0: state = { 'encoder_state_dict': self.encoder.state_dict(), 'memory_state_dict': self.memory.state_dict(), 'decoder_state_dict': self.decoder.state_dict(), 'acc': acc, 'session': session, 'epoch': epoch, } file_path = osp.join( self.args.save_dir, self.dataset, 'incr', self.args.encoder, 'session' + str(session), 'ckp_ep' + str(epoch) + '.ckp') if epoch == 2: pass # save_checkpoint(state, False, file_path) else: save_checkpoint(state, is_best, file_path) if is_best: best_acc = acc best_epoch = epoch best_state = copy.deepcopy(state) print('==> Test best accuracy {:.2%}, achieved at epoch {}'.format( best_acc, best_epoch)) torch.cuda.empty_cache() # Load best checkpoint self.encoder.load_state_dict(best_state['encoder_state_dict']) self.memory.load_state_dict(best_state['memory_state_dict']) self.decoder.load_state_dict(best_state['decoder_state_dict']) elapsed = str(datetime.timedelta(seconds=round(time.time() - start_time))) train_time = str(datetime.timedelta(seconds=train_time)) print('Session {} finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}'.format( session, elapsed, train_time)) # print("==========\nArgs:{}\n==========".format(self.args)) memory_vals = self.memory.m_vals.cpu().numpy() memory_vals_counter = Counter(memory_vals) print('memory val:', len(memory_vals_counter), memory_vals_counter.most_common())
def main(): torch.manual_seed(1) os.environ['CUDA_VISIBLE_DEVICES'] = config.gpu_devices use_gpu = torch.cuda.is_available() sys.stdout = Logger(config.save_dir, config.checkpoint_suffix, config.evaluate) print("\n==========\nArgs:") config.print_parameter() print("==========\n") if use_gpu: print("Currently using GPU {}".format(config.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(1) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(config.dataset)) dataset = data_manager.init_imgreid_dataset(name=config.dataset, root=config.data_root) transform_train = T.Compose([ T.Random2DTranslation(config.height, config.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=data_mean, std=data_std), ]) transform_test = T.Compose([ T.Resize((config.height, config.width)), T.ToTensor(), T.Normalize(mean=data_mean, std=data_std), ]) pin_memory = True if use_gpu else False # train_batch_sampler = BalancedBatchSampler(dataset.train, n_classes=8, n_samples=8) # train_batch_sampler = CCLBatchSampler(dataset.train, n_classes=n_classes, n_samples=n_samples) # train_batch_sampler = CCLBatchSamplerV2(dataset.train, n_classes=n_classes, pos_samp_cnt=pos_samp_cnt, # neg_samp_cnt=neg_samp_cnt, each_cls_max_cnt=each_cls_max_cnt) train_batch_sampler = ClassSampler(dataset.train, sample_cls_cnt=config.sample_cls_cnt, each_cls_cnt=config.each_cls_cnt) # trainloader = DataLoader( # ImageDataset(dataset.train, transform=transform_train), # batch_sampler=train_batch_sampler, batch_size=args.train_batch, # shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True # ) trainloader = DataLoader(ImageDatasetWCL(dataset, data_type='train', merge_h=256, merge_w=256, mean_std=[data_mean, data_std]), batch_sampler=train_batch_sampler, num_workers=config.workers, pin_memory=pin_memory) queryloader = DataLoader( ImageDatasetWCL(dataset.query, data_type='query', merge_h=256, merge_w=256, mean_std=[data_mean, data_std]), batch_size=config.test_batch, shuffle=False, num_workers=config.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( ImageDatasetWCL(dataset.gallery, data_type='gallery', merge_h=256, merge_w=256, mean_std=[data_mean, data_std]), batch_size=config.test_batch, shuffle=False, num_workers=config.workers, pin_memory=pin_memory, drop_last=False, ) if config.dataset == 'vehicleid': train_query_loader = None train_gallery_loader = None else: train_query_loader = DataLoader( ImageDatasetWCL(dataset.train_query, data_type='train_query', merge_h=256, merge_w=256, mean_std=[data_mean, data_std]), batch_size=config.test_batch, shuffle=False, num_workers=config.workers, pin_memory=pin_memory, drop_last=False, ) train_gallery_loader = DataLoader( ImageDatasetWCL(dataset.train_gallery, data_type='train_gallery', merge_h=256, merge_w=256, mean_std=[data_mean, data_std]), batch_size=config.test_batch, shuffle=False, num_workers=config.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(config.arch)) model = init_model(name=config.arch, num_classes=dataset.num_train_pids, loss_type=config.loss_type) print("Model size: {:.3f} M".format(count_num_param(model))) if config.loss_type == 'xent': criterion = [nn.CrossEntropyLoss(), nn.CrossEntropyLoss()] elif config.loss_type == 'xent_triplet': criterion = XentTripletLoss( margin=config.margin, triplet_selector=RandomNegativeTripletSelector( margin=config.margin), each_cls_cnt=config.each_cls_cnt, n_class=config.sample_cls_cnt) elif config.loss_type == 'xent_tripletv2': criterion = XentTripletLossV2( margin=config.margin, triplet_selector=RandomNegativeTripletSelectorV2( margin=config.margin), each_cls_cnt=config.each_cls_cnt, n_class=config.sample_cls_cnt) # criterion = XentTripletLossV2(margin=0.04, triplet_selector=RandomNegativeTripletSelectorV2(margin=0.04), # each_cls_cnt=config.each_cls_cnt, n_class=config.sample_cls_cnt) # criterion = XentGroupTripletLossV2(margin=0.8, triplet_selector=AllTripletSelector(margin=0.8), # each_cls_cnt=config.each_cls_cnt, n_class=config.sample_cls_cnt) else: raise KeyError("Unsupported loss: {}".format(config.loss_type)) optimizer = init_optim(config.optim, model.parameters(), config.lr, config.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=config.stepsize, gamma=config.gamma) if config.resume is not None: if check_isfile(config.resume): checkpoint = torch.load(config.resume) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) config.start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] if 'mAP' in checkpoint: mAP = checkpoint['mAP'] else: mAP = 0 print("Loaded checkpoint from '{}'".format(config.resume)) print("- start_epoch: {}\n- rank1: {}\n- mAP: {}".format( config.start_epoch, rank1, mAP)) if use_gpu: model = nn.DataParallel(model).cuda() if config.evaluate: print("Evaluate only") test_model(model, queryloader, galleryloader, train_query_loader, train_gallery_loader, use_gpu, config.test_batch, config.loss_type, config.euclidean_distance_loss) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_map = 0 best_epoch = 0 for epoch in range(config.start_epoch, config.max_epoch): print("==> Start training") start_train_time = time.time() scheduler.step() print('epoch:', epoch, 'lr:', scheduler.get_lr()) train(epoch, model, criterion, optimizer, trainloader, config.loss_type, config.print_freq) train_time += round(time.time() - start_train_time) if epoch >= config.start_eval and config.eval_step > 0 and epoch % config.eval_step == 0 \ or epoch == config.max_epoch: print("==> Test") rank1, mAP = test_model(model, queryloader, galleryloader, train_query_loader, train_gallery_loader, use_gpu, config.test_batch, config.loss_type, config.euclidean_distance_loss) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_map = mAP 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, 'mAP': mAP, 'epoch': epoch + 1, }, is_best, use_gpu_suo=False, fpath=osp.join( config.save_dir, 'checkpoint_ep' + str(epoch + 1) + config.checkpoint_suffix + '.pth.tar')) print("==> Best Rank-1 {:.2%}, mAP {:.2%}, achieved at epoch {}".format( best_rank1, best_map, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_vidreid_dataset(root=args.root, name=args.dataset) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False # decompose tracklets into images for image-based training new_train = [] for img_paths, pid, camid in dataset.train: for img_path in img_paths: new_train.append((img_path, pid, camid)) trainloader = DataLoader( ImageDataset(new_train, transform=transform_train), sampler=RandomIdentitySampler(new_train, num_instances=args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( VideoDataset(dataset.query, seq_len=args.seq_len, sample='evenly', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( VideoDataset(dataset.gallery, seq_len=args.seq_len, sample='evenly', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(model))) 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) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if args.load_weights: # load pretrained weights but ignore layers that don't match in size print("Loading pretrained weights from '{}'".format(args.load_weights)) checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()} model_dict.update(pretrain_dict) model.load_state_dict(model_dict) if args.resume: checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, args.pool, use_gpu) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch+1) > args.start_eval and args.eval_step > 0 and (epoch+1) % args.eval_step == 0 or (epoch+1) == args.max_epoch: print("==> Test") rank1 = test(model, queryloader, galleryloader, args.pool, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint({ 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch+1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time))