def __init__(self, load=True, width=256, height=256): self.writer = SummaryWriter('runs/SCL_1') self.spatial_features_size = 20 self.load = load self.mod = False self.reverse_margins = False self.dataset = SCL_dataset(width, height, self.mod, self.reverse_margins) self.model = SCL_model(self.spatial_features_size, width, height).to(device) self.descriptor_model = Descriptor_net(self.spatial_features_size, width, height).to(device) self.contrastive_criterion = TripletLoss(self.mod, self.reverse_margins) # self.pixel_criterion=PixelwiseLoss(width, height) self.pixel_criterion = PixelTripletLoss(width, height, self.mod) self.load_from = "models/SCL.pth" self.save_to = "models/SCL.pth" self.d_load_from = "models/d_SCL.pth" self.d_save_to = "models/d_SCL.pth" if self.load: self.model.load_state_dict( torch.load(self.load_from, map_location=device)) self.model.eval() self.delayed_model = SCL_model(self.spatial_features_size, width, height).to(device) self.delayed_model.load_state_dict( torch.load(self.load_from, map_location=device)) self.delayed_model.eval() self.descriptor_model.load_state_dict( torch.load(self.d_load_from, map_location=device)) self.descriptor_model.eval() self.d_delayed_model = Descriptor_net(self.spatial_features_size, width, height).to(device) self.d_delayed_model.load_state_dict( torch.load(self.d_load_from, map_location=device)) self.d_delayed_model.eval() self.tau = 0.001 self.optimizer = optim.Adam(self.model.parameters(), lr=1e-4, weight_decay=1e-4) self.descriptor_optimizer = optim.Adam( self.descriptor_model.parameters(), lr=1e-4, weight_decay=1e-4) self.max_iter = 3500 self.dataloader = torch.utils.data.DataLoader(self.dataset, 1, shuffle=True, pin_memory=device) self.descriptor_dataloader = torch.utils.data.DataLoader( self.dataset, 1, shuffle=True, pin_memory=device) self.images = self.dataset.frames[0] self.embedding = torch.empty((len(self.images), 32)) self.cost_counter = 0
def train_model(triplet_train_loader, triplet_val_loader, test_fun, margin, net_output, model_directory): target_embedding_net = TargetEmbeddingNet(300, net_output) triple_model = TripletNet(target_embedding_net, target_embedding_net) if cuda: triple_model.cuda() loss_fn = TripletLoss(margin) lr = 1e-4 optimizer = optim.Adam(triple_model.parameters(), lr=lr) scheduler = lr_scheduler.StepLR(optimizer, 1000, gamma=0.5, last_epoch=-1) n_epochs = 20 log_interval = 100 train_losses, val_losses, metrices = fit(triplet_train_loader, triplet_val_loader, triple_model, loss_fn, optimizer, scheduler, n_epochs, cuda, log_interval, callback_test=test_fun, keep_checkpoint_max=10, model_dir=model_directory) return metrices
def net_config(image, label, model, args, is_train): assert args.model in model_list, "{} is not in lists: {}".format( args.model, model_list) out = model.net(input=image, embedding_size=args.embedding_size) if not is_train: return None, out if args.loss_name == "triplet": metricloss = TripletLoss(margin=args.margin, ) elif args.loss_name == "quadruplet": metricloss = QuadrupletLoss( train_batch_size=args.train_batch_size, samples_each_class=args.samples_each_class, margin=args.margin, ) elif args.loss_name == "eml": metricloss = EmlLoss( train_batch_size=args.train_batch_size, samples_each_class=args.samples_each_class, ) elif args.loss_name == "npairs": metricloss = NpairsLoss( train_batch_size=args.train_batch_size, samples_each_class=args.samples_each_class, reg_lambda=args.npairs_reg_lambda, ) cost = metricloss.loss(out, label) avg_cost = fluid.layers.mean(x=cost) return avg_cost, out
def run(train_loader, test_pairs, options): # Construct model cuda = options["device_id"] != -1 if options["use_pair_feature"]: point_net = PointNet(options) pair_net = PairNetWithPairFeatures(options, point_net) model = TripletNetWithPairFeatures(pair_net) else: point_net = PointNet(options) pair_net = PairNet(options, point_net) model = TripletNet(pair_net) if cuda: model.cuda() print(model) optimizer = torch.optim.Adam(model.parameters(), lr=options["lr"], weight_decay=options["weight_decay"]) loss_fn = TripletLoss(options["margin"]) best_overall_metric = 0 # a single value best_metrics = None # a dictionary best_epoch = 0 save_model(model, options["save_dir"], 'best', 0) # save the initial first model for epoch in range(options["epochs"]): epoch_loss, non_zero_triplet_ratio = train_triplet_epoch(train_loader, model, loss_fn, optimizer, cuda, use_pair_feature=options["use_pair_feature"]) print("Epoch: {}, train-loss: {:06.4f}, non_zero_triplet_ratio: {:06.4f}, ".format(epoch, epoch_loss, non_zero_triplet_ratio)) if epoch % options["eval_epoch"] == 0 and epoch != 0: if options["use_pair_feature"]: prediction_score = pair_prediction_with_pair_feature(model.pair_net, test_pairs, pair_features, cuda, options, batch_size=10000) else: prediction_score = pair_prediction(model.pair_net, test_pairs, cuda, options, batch_size=10000) test_triplets = [] for term_pair, score in zip(test_pairs, prediction_score): test_triplets.append((term_pair[0], term_pair[1], -1.0 * score)) metrics = evaluation_main(test_triplets) if metrics["all"] >= best_overall_metric: best_overall_metric =metrics["all"] best_epoch = epoch best_metrics = metrics save_model(model, options["save_dir"], 'best', epoch) # save the initial first model return best_overall_metric, best_epoch, best_metrics
def __init__(self, settings, input_shape=(1, 28, 28), embedding_size=10, margin_loss=0.2): super(EmbeddingNet, self).__init__() in_channels = input_shape[0] sequence = [nn.Conv2d(in_channels, 128, kernel_size=7, padding=1), # (28, 28) nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(128, 128, kernel_size=3, padding=1), # (11, 11) nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(128, 256, kernel_size=3, padding=1), # (5, 5) nn.ReLU(inplace=True), Flatten(), # (4, 4) nn.Linear(6 * 6 * 256, 4096), nn.ReLU(inplace=True), nn.Linear(4096, embedding_size), nn.LayerNorm(embedding_size, elementwise_affine=False) ] self.model = nn.Sequential(*sequence) self.model.apply(weight_init) self.loss_fun = TripletLoss(margin_loss).triplet_loss self.cum_loss = 0.0 self.scheduler, self.optimizer = self.init_optimizer(settings) self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
def parse_parameters(params, optimizer, lr, loss, schedule=None): optimizer = optimizer.lower() if 'adam' in optimizer: optimizer = optim.Adam(lr=lr, params=params, weight_decay=1e-5) elif 'sgd' in optimizer: optimizer = optim.SGD(lr=lr, momentum=0.9, params=params) else: print('optimizer {} not yet supported'.format(optimizer)) scheduler = None if schedule is not None: scheduler = optim.lr_scheduler.ExponentialLR(optimizer, math.exp(schedule)) loss = loss.lower() if 'cont' in loss: criterion = ContrastiveLoss() elif 'trip' in loss: criterion = TripletLoss() else: print('loss {} not supported'.format(loss)) return (optimizer, criterion, scheduler)
def train(train_path, labels, boxes, output_dir, num_epochs, hard, verbose): df_train, df_val = _get_toy_dataset(labels, boxes) if verbose: logging.info("Train size: {}, validation size: {}".format( len(df_train), len(df_val))) sampler = None if hard: sampler = BalancedBatchSampler(df_train, n_classes=4, n_samples=4) train_dl, single_train_dl, val_dl = get_dataloaders( df_train, df_val, train_path, sampler) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if verbose: logging.info("Using device {}".format(device)) net = TripletNetwork(embedding_size=128).to(device) criterion = TripletLoss() selector = None if hard: selector = HardBatchTripletSelector() optimizer = optim.Adam(net.parameters(), lr=1e-4) net, history = _train( model=net, optimizer=optimizer, criterion=criterion, train_dataloader=train_dl, single_train_dataloader=single_train_dl, val_dataloader=val_dl, num_epochs=num_epochs, save_path=output_dir, device=device, selector=selector, ) _plot_history(history)
negative, negative_labels = load_data(NEGATIVE_DIR, n_class=NUM_CLASS, cls=5, img_size=IMAGE_SIZE) print("-" * 30) print("Data loaded!!") print("-" * 30) # 3 models extracting the pattern feature triplet_mdoel = TripletNet(CNNEmbeddingNetL2()).cuda() # pattern_model = TripletNet( # torch.hub.load('pytorch/vision:v0.9.0', 'mobilenet_v2', pretrained=True) # ).cuda() criterion = TripletLoss(margin=MARGIN) # optimizer = optim.Adam(triplet_model.parameters(), lr=LR) optimizer = optim.SGD(triplet_mdoel.embedding_net.parameters(), lr=LR, momentum=0.9) """Training Phase""" """ Stage 1: Train the Embedding Network""" triplet_model = train_triplet(triplet_mdoel, criterion, optimizer, training, training_labels, negative, n_epoch=N_EPOCH, batch_size=TRAINING_BATCH_SIZE)
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)))
# init the logger init_logger(args) # data loading dataset, trainloader, queryloader, galleryloader = init_data_loaders(args) num_train_pids = dataset.num_train_pids # init model model = init_model(args, num_train_pids) if use_gpu: model = nn.DataParallel(model).cuda() vis = utils.get_visdom_for_current_run(args.save_dir, args.prefix + '_stage1_training') # init objective functions criterion_xent = CrossEntropyLabelSmooth(num_classes=num_train_pids, use_gpu=use_gpu) criterion_htri = TripletLoss(margin=args.margin) # init optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) if args.stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch # if only evaluation was needed if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, use_gpu, args) exit(0) start_time = time.time()
def main(): import torch from torch.optim import lr_scheduler import torch.optim as optim from torch.autograd import Variable from trainer import fit import numpy as np cuda = torch.cuda.is_available() # Training settings parser = argparse.ArgumentParser(description='PyTorch MNIST Example') parser.add_argument('--batch-size', type=int, default=100, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('--test-batch-size', type=int, default=100, metavar='N', help='input batch size for testing (default: 1000)') parser.add_argument('--epochs', type=int, default=100, metavar='N', help='number of epochs to train (default: 10)') parser.add_argument('--lr', type=float, default=0.001, metavar='LR', help='learning rate (default: 0.01)') parser.add_argument('--momentum', type=float, default=0.5, metavar='M', help='SGD momentum (default: 0.5)') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)') parser.add_argument( '--log-interval', type=int, default=10, metavar='N', help='how many batches to wait before logging training status') parser.add_argument('--save-model', action='store_true', default=True, help='For Saving the current Model') args = parser.parse_args() use_cuda = not args.no_cuda and torch.cuda.is_available() torch.manual_seed(args.seed) np.random.seed(args.seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False device = torch.device("cuda" if use_cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} from datetime import datetime import os loging = False ismultitask = False loso = False if (args.save_model): model_save_path = 'model/triplet/' if not os.path.isdir(model_save_path): os.makedirs(model_save_path) if loging: fname = model_save_path + datetime.today().strftime( "%m_%d_%H_%M") + ".txt" f = open(fname, 'w') x_data, y_data = load_smt() # nonbciilli = np.s_[0,1,2,4,5,8,16,17,18,20,21,27,28,29,30,32,35,36,38,42,43,44,51] valtype = 'sess' if valtype == 'loso': for subj in range(0, 54): model = Deep4CNN(ismult=ismultitask).to(device) #model.load_state_dict(torch.load(model_save_path+ "J_" + str(subj) + 'basecnn.pt')) optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) optimizer_fine = optim.SGD(model.parameters(), lr=0.005, momentum=args.momentum) dataset_train = GigaDataset(x=x_data, y=y_data, valtype=valtype, istrain=True, sess=1, subj=subj) train_loader = torch.utils.data.DataLoader( dataset_train, batch_size=args.batch_size, shuffle=True, **kwargs) dataset_test = GigaDataset(x=x_data, y=y_data, valtype=valtype, istrain=False, sess=2, subj=subj) test_loader = torch.utils.data.DataLoader( dataset_test, batch_size=args.batch_size, shuffle=False, **kwargs) # dataset_fine = GigaDataset_LOSO(x=x_data, y=y_data, fine=True, istrain=True, sess=2, subj=subj) # fine_loader = torch.utils.data.DataLoader(dataset_fine, batch_size=args.batch_size, shuffle=True, **kwargs) for epoch in range(1, args.epochs + 1): train(args, model, device, train_loader, optimizer, epoch) print("joint-train") #LOSO joint training j_loss, j_score = eval(args, model, device, test_loader) if epoch > 30: if (args.save_model): torch.save( model.state_dict(), model_save_path + "model_" + str(subj) + "_" + str(epoch) + '.pt') # #fine tuning # for epoch in range(1, 10): # train_mt(args, model, device, fine_loader, optimizer_fine, epoch) # # print("fine-tuning") # f_loss, f_score = eval(args, model, device, test_loader) if (args.save_model): torch.save(model.state_dict(), model_save_path + "F_" + str(subj) + 'basecnn.pt') if loging: f = open(fname, 'a') f.write( str(subj) + " " + "jl : " + str(j_loss) + " " + str(j_score) + '\n') f.close() elif valtype == 'sess': from networks import EmbeddingDeep4CNN, TripletNet, FineShallowCNN, EmbeddingDeepCNN from losses import TripletLoss margin = 1 embedding_net = EmbeddingDeep4CNN() print(embedding_net) model = TripletNet(embedding_net) if cuda: model.cuda() loss_fn = TripletLoss(margin) lr = 1e-3 #optimizer = optim.Adam(model.parameters(), lr=lr) n_epochs = 5 #%% log_interval = 10 if n_epochs == 0: pass #model.load_state_dict(torch.load('triplet_deep4_1000_2.pt')) else: #트리플렛넷 학습 # For classification dataset_train = GigaDataset(x=x_data, y=y_data, valtype=valtype, istrain=True, sess=1) train_loader = torch.utils.data.DataLoader( dataset_train, batch_size=args.batch_size, shuffle=True, **kwargs) dataset_test = GigaDataset(x=x_data, y=y_data, valtype=valtype, istrain=False, sess=2, subj=-1) test_loader = torch.utils.data.DataLoader( dataset_test, batch_size=args.batch_size, shuffle=False, **kwargs) triplet_dataset_train = TripletGiga(x=x_data, y=y_data, valtype=valtype, istrain=True, sess=1) triplet_train_loader = torch.utils.data.DataLoader( triplet_dataset_train, batch_size=args.batch_size, shuffle=True, **kwargs) triplet_dataset_test = TripletGiga(x=x_data, y=y_data, valtype=valtype, istrain=False, sess=2, subj=-1) triplet_test_loader = torch.utils.data.DataLoader( triplet_dataset_test, batch_size=args.batch_size, shuffle=False, **kwargs) optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) scheduler = lr_scheduler.StepLR(optimizer, 8, gamma=1, last_epoch=-1) from trainer import fit fit(triplet_train_loader, triplet_test_loader, model, loss_fn, optimizer, scheduler, n_epochs, cuda, log_interval) #%% train_embeddings_tl, train_labels_tl = extract_embeddings( train_loader, embedding_net, 1000) # plot_embeddings(train_embeddings_tl, train_labels_tl) val_embeddings_tl, val_labels_tl = extract_embeddings( test_loader, embedding_net, 1000) # plot_embeddings(val_embeddings_tl, val_labels_tl) # # from sklearn.pipeline import Pipeline from sklearn.discriminant_analysis import LinearDiscriminantAnalysis from sklearn.model_selection import ShuffleSplit, cross_val_score lda = LinearDiscriminantAnalysis() lda.fit(train_embeddings_tl, train_labels_tl) print(lda.score(val_embeddings_tl, val_labels_tl)) # from torchvision import datasets, models, transforms # temp = model.embedding_net.children() # newmodel = torch.nn.Sequential(*(list(model.embedding_net.children())[:])) # for param in model.embedding_net.parameters(): # param.requires_grad = True #newembedding_net = torch.nn.Sequential(*(list(model.embedding_net.children())[:])) # from sklearn.manifold import TSNE tsne = TSNE(n_components=2, perplexity=30) train_tsne = tsne.fit_transform(val_embeddings_tl) plot_embeddings(train_tsne, val_labels_tl) for param in model.embedding_net.parameters(): param.requires_grad = True #embedding_net2 = EmbeddingDeep4CNN() newmodel = nn.Sequential(model.embedding_net, nn.Linear(1000, 2), nn.LogSoftmax(dim=1)).to(device) print(newmodel) #newmodel.fc_lr = nn.Linear(1000,2) newmodel.to(device) optimizer = optim.SGD(newmodel.parameters(), lr=0.01, momentum=0.9) #optimizer = optim.Adam(newmodel.parameters()) for epoch in range(1, 20): train(args, newmodel, device, train_loader, optimizer, epoch) j_loss, j_score = eval(args, newmodel, device, test_loader) if args.save_model: torch.save(model.state_dict(), 'triplet_deep4_1000_2.pt')
def testseq(dataset_name, use_gpu): dataset_root = './video2img/track1_sct_img_test_big/' dataset = Graph_data_manager.AICityTrack2(root=dataset_root) width = 224 height = 224 transform_train = T.Compose([ T.Random2DTranslation(height, 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((height, 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 seq_len = 4 num_instance = 4 train_batch = 32 test_batch = 1 queryloader = DataLoader( VideoDataset(dataset.query, seq_len=seq_len, sample='dense', transform=transform_test), batch_size=test_batch, shuffle=False, num_workers=4, pin_memory=pin_memory, drop_last=False, ) arch = "resnet50ta" pretrained_model = "./log/track12_ta224_checkpoint_ep500.pth.tar" start_epoch = 0 print("Initializing model: {}".format(arch)) dataset.num_train_pids = 517 if arch == 'resnet503d': model = resnet3d.resnet50(num_classes=dataset.num_train_pids, sample_width=width, sample_height=height, sample_duration=seq_len) if not os.path.exists(pretrained_model): raise IOError( "Can't find pretrained model: {}".format(pretrained_model)) print("Loading checkpoint from '{}'".format(pretrained_model)) checkpoint = torch.load(pretrained_model) state_dict = {} for key in checkpoint['state_dict']: if 'fc' in key: continue state_dict[key.partition("module.") [2]] = checkpoint['state_dict'][key] model.load_state_dict(state_dict, strict=False) else: if not os.path.exists(pretrained_model): model = models.init_model(name=arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) else: model = models.init_model(name=arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) checkpoint = torch.load(pretrained_model) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] + 1 print("Loaded checkpoint from '{}'".format(pretrained_model)) print("- start_epoch: {}\n- rank1: {}".format( start_epoch, checkpoint['rank1'])) 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=0.3) lr = 0.0003 gamma = 0.1 stepsize = 200 weight_decay = 5e-04 optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) if stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=stepsize, gamma=gamma) start_epoch = start_epoch if use_gpu: model = nn.DataParallel(model).cuda() test(model, queryloader, 'avg', use_gpu, dataset, -1, meta_data_tab=None)
def main(): args.save_dir = args.save_dir + '/' + args.arch 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 # add data to save_dir args.save_dir = args.save_dir + '_' + args.dataset + '_combined_multisteplr11' if args.pretrained_model is not None: args.save_dir = os.path.dirname(args.pretrained_model) if not osp.exists(args.save_dir): os.makedirs(args.save_dir) log_name = 'test.log' if args.evaluate else 'train.log' log_name += time.strftime('-%Y-%m-%d-%H-%M-%S') sys.stdout = Logger(osp.join(args.save_dir, log_name)) 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_dataset(name=args.dataset) print("Train Transforms: \n\ Random2DTranslation, \n\ RandomHorizontalFlip, \n\ ToTensor, \n\ normalize\ ") transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), # T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # T.RandomErasing(p=0.5, scale=(0.02, 0.4), ratio=(0.3, 3.3), value=[0.485, 0.456, 0.406]) ]) 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( VideoDataset(dataset.train, seq_len=args.seq_len, sample=args.data_selection, 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( VideoDataset(dataset.query, seq_len=args.seq_len, sample='dense', 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='dense', 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, seq_len=args.seq_len) # pretrained model loading if args.pretrained_model is not None: if not os.path.exists(args.pretrained_model): raise IOError("Can't find pretrained model: {}".format( args.pretrained_model)) print("Loading checkpoint from '{}'".format(args.pretrained_model)) pretrained_state = torch.load(args.pretrained_model)['state_dict'] print(len(pretrained_state), ' keys in pretrained model') current_model_state = model.state_dict() pretrained_state = {key: val for key, val in pretrained_state.items() if key in current_model_state and val.size() == current_model_state[key].size()} print(len(pretrained_state), ' keys in pretrained model are available in current model') current_model_state.update(pretrained_state) model.load_state_dict(current_model_state) print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters())/1000000.0)) if use_gpu: model = nn.DataParallel(model).cuda() criterion_xent = CrossEntropyLabelSmooth( num_classes=dataset.num_train_pids, use_gpu=use_gpu) criterion_htri = TripletLoss(margin=args.margin) optimizer = torch.optim.Adam( model.parameters(), lr=args.lr, weight_decay=args.weight_decay) if args.stepsize > 0: scheduler = lr_scheduler.StepLR( optimizer, step_size=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, use_gpu) return start_time = time.time() best_rank1 = -np.inf is_first_time = True for epoch in range(start_epoch, args.max_epoch): eta_seconds = (time.time() - start_time) * (args.max_epoch - epoch) / max(epoch, 1) eta_str = str(datetime.timedelta(seconds=int(eta_seconds))) print("==> Epoch {}/{} \teta {}".format(epoch+1, args.max_epoch, eta_str)) train(model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) if args.stepsize > 0: scheduler.step() rank1 = 'NA' mAP = 'NA' is_best = False if args.eval_step > 0 and (epoch+1) % args.eval_step == 0 or (epoch+1) == args.max_epoch: print("==> Test") rank1, mAP = test(model, queryloader, galleryloader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 # save the model as required if (epoch+1) % args.save_step == 0: 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, args.save_prefix, 'model' + '.pth.tar-' + str(epoch+1))) is_first_time = False if not is_first_time: utils.disable_all_print_once() elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
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_vid_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: {:.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: 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, 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) if args.stepsize > 0: scheduler.step() if 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))
# determine the epoch of pretrained model epoch = int((opt.pretrained_model.split("/")[-1]).split("-")[-1]) # if only evaluation is required if opt.evaluate: print("-- evaluate only") test(epoch, model) exit(0) # Define optimizer and loss function person_id_criterion = CrossEntropyLabelSmooth( train_dataset.number_classes(), use_gpu=opt.cuda) attribute_criterion = AttributeCriterion(attribute_choices, CrossEntropyLabelSmooth) triplet_criterion = TripletLoss(opt.margin) optimizer = optim.Adam(model.parameters(), lr=opt.lr, weight_decay=5e-4) # Default lr = 3e-4 print("Using triplet loss = ", triplet_criterion) print("Using person_id = ", person_id_criterion) print("Using Attribute loss = ", attribute_criterion) print("Optimizer = ", optimizer) # scheduler creation lr_scheduler, num_epochs = create_scheduler(opt, optimizer) if epoch > 0: lr_scheduler.step(epoch) print("Scheduled epochs: ", num_epochs) print("learning rates ",
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 main(): # 为了看看repo提供的model.pth.tar-9在validation集的mAp和rank-1 # 我自己训练的tar-9只有mAP: 15.1%; Rank-1: 23.3% ,不知道为什么 # 更改args.load_weights = '/model/caohw9/track3_model/model.pth.tar-9' global args print(args) set_random_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 sys.stdout = Logger(osp.join(args.save_dir, "log.txt")) if use_gpu: print('Currently using GPU {}'.format(args.gpu_devices)) cudnn.benchmark = True else: warnings.warn( 'Currently using CPU, however, GPU is highly recommended') # 初始化loader print('Initializing image data manager') dm = ImageDataManager(use_gpu, **trainset_kwargs(args)) trainloader, testloader_dict = dm.return_dataloaders( ) #trainloader用于训练,testloader_dict包含['query']和['gallery']2个loader print('suffessfully initialize loaders!') # 初始化模型 print('Initializing model: {}'.format( args.arch)) #args.arch default='resnet101' model = models.init_model(name=args.arch, num_classes=dm.num_train_pids, loss={'xent', 'htri'}, pretrained=not args.no_pretrained, use_gpu=use_gpu) print('Model size: {:.3f} M'.format(count_num_param(model))) # 加载预训练参数 if args.load_weights and check_isfile(args.load_weights): load_pretrained_weights(model, args.load_weights) #加载训练过的模型后,先看看validation print('=> Validation') print('Evaluating {} ...'.format( args.test_set)) #args.test_set应该是指的validation set? queryloader = testloader_dict['query'] galleryloader = testloader_dict['test'] model = nn.DataParallel(model).cuda() if use_gpu else model rank1 = test(model, queryloader, galleryloader, use_gpu) #validation! # 多GPU训练 else: model = nn.DataParallel(model).cuda() if use_gpu else model # 定义loss,optimizer, lr_scheduler criterion_xent = CrossEntropyLoss(num_classes=dm.num_train_pids, use_gpu=use_gpu, label_smooth=args.label_smooth) criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optimizer(model, **optimizer_kwargs(args)) scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs(args)) # 是否是resume训练 if args.resume and check_isfile(args.resume): args.start_epoch = resume_from_checkpoint( args.resume, model, optimizer=optimizer) #获取中断时刻的epoch数 # 开始训练! time_start = time.time() print('=> Start training') for epoch in range(args.start_epoch, args.max_epoch): train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) #训练 scheduler.step() #更新lr # 当epoch数超过args.start_eval,每隔一定频率args.eval_freq,或者达到最后一个epoch,进行validation+存储checkpoint if (epoch + 1) > args.start_eval and args.eval_freq > 0 and ( epoch + 1) % args.eval_freq == 0 or (epoch + 1) == args.max_epoch: print('=> Validation') print('Evaluating {} ...'.format( args.test_set)) #args.test_set应该是指的validation set? queryloader = testloader_dict['query'] galleryloader = testloader_dict['test'] rank1 = test(model, queryloader, galleryloader, use_gpu) #validation! save_checkpoint( { 'state_dict': model.state_dict(), #模型的状态字典 'rank1': rank1, 'epoch': epoch + 1, 'arch': args.arch, #default='resnet101' 'optimizer': optimizer.state_dict( ), #优化器对象的状态字典,包含优化器的状态和超参数(如lr, momentum,weight_decay等) }, args.save_dir) #validation同时保存checkpoint # 训练结束! elapsed = round(time.time() - time_start) #持续时间 elapsed = str(datetime.timedelta(seconds=elapsed)) print('Elapsed {}'.format(elapsed))
def run_experiment(sim_type, network_type, dataset, mus, latent_dims, batch_size=64, keys=None, composite_labels=None, margin=1, dropout_val=0.2, num_iters=20000, normalize=True, feature_extractor='hog'): """Overarching function that runs an experiment. Calls _experiment for each combination of mu and latent dim. Args: sim_type (str): Simulation type. Accepts either "pairwise" or "triplet" network_type (str): Network type. Accepts either "embenc" or "vae" dataset (str): Dataset name. Accepts "mnist", "cifar10" and "imagenet" mus (list): list of mu values for which experiment is to be run. latent_dims (list): list of latent dimensions for which experiment is to be run. batch_size (int, optional): batch size from which pairs/triples are generated. Defaults to 64. keys (dict, optional): Mapping from numeric labels to string label names. Defaults to None. composite_labels (dict, optional): Maps current labels to composite labels. Defaults to None. margin (int, optional): Margin to be used for metric loss. Defaults to 1. dropout_val (float, optional): dropout value for networks. Defaults to 0.2. num_iters (int, optional): number of iterations. Defaults to 20000. Returns: dict: mapping from (mu, latent_dim) to results for _experiment on those params. """ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print("loading data...", end='') trainX, trainY, testX, testY, trainLabels, testLabels = load_data( dataset, composite_labels=composite_labels, normalize=normalize, feature_extractor=feature_extractor) print("done.") # instantiate the generator and the networks based on type. print("initializing data generators and losses...", end='') input_dim = trainX.shape[-1] if sim_type == 'pairwise': datagen = PairGenerator(trainX, trainY, batch_size) metric_loss = ContrastiveLoss( margin) if network_type != 'vae' else PairwiseVAELoss(margin) elif sim_type == 'triplet': datagen = TripletGenerator(trainX, trainY, batch_size) metric_loss = TripletLoss( margin) if network_type != 'vae' else TripletVAELoss(margin) reconstruction_loss = reconstruction_error print("done.") all_results = {} for mu in mus: for latent_dim in latent_dims: network = initialize_network(sim_type, network_type, input_dim, latent_dim, dropout_val, device=device) results = _experiment(network, latent_dim, datagen, metric_loss, reconstruction_loss, trainX, trainY, testX, testY, keys, mu, dataset, trainLabels, testLabels, num_iters=num_iters, batch_size=batch_size, device=device) all_results[(mu, latent_dim)] = results return all_results
class SCL_trainer: def __init__(self, load=True, width=256, height=256): self.writer = SummaryWriter('runs/SCL_1') self.spatial_features_size = 20 self.load = load self.mod = False self.reverse_margins = False self.dataset = SCL_dataset(width, height, self.mod, self.reverse_margins) self.model = SCL_model(self.spatial_features_size, width, height).to(device) self.descriptor_model = Descriptor_net(self.spatial_features_size, width, height).to(device) self.contrastive_criterion = TripletLoss(self.mod, self.reverse_margins) # self.pixel_criterion=PixelwiseLoss(width, height) self.pixel_criterion = PixelTripletLoss(width, height, self.mod) self.load_from = "models/SCL.pth" self.save_to = "models/SCL.pth" self.d_load_from = "models/d_SCL.pth" self.d_save_to = "models/d_SCL.pth" if self.load: self.model.load_state_dict( torch.load(self.load_from, map_location=device)) self.model.eval() self.delayed_model = SCL_model(self.spatial_features_size, width, height).to(device) self.delayed_model.load_state_dict( torch.load(self.load_from, map_location=device)) self.delayed_model.eval() self.descriptor_model.load_state_dict( torch.load(self.d_load_from, map_location=device)) self.descriptor_model.eval() self.d_delayed_model = Descriptor_net(self.spatial_features_size, width, height).to(device) self.d_delayed_model.load_state_dict( torch.load(self.d_load_from, map_location=device)) self.d_delayed_model.eval() self.tau = 0.001 self.optimizer = optim.Adam(self.model.parameters(), lr=1e-4, weight_decay=1e-4) self.descriptor_optimizer = optim.Adam( self.descriptor_model.parameters(), lr=1e-4, weight_decay=1e-4) self.max_iter = 3500 self.dataloader = torch.utils.data.DataLoader(self.dataset, 1, shuffle=True, pin_memory=device) self.descriptor_dataloader = torch.utils.data.DataLoader( self.dataset, 1, shuffle=True, pin_memory=device) self.images = self.dataset.frames[0] self.embedding = torch.empty((len(self.images), 32)) self.cost_counter = 0 def run_training(self, iterations=100): self.dataset.count = 0 i = 0 for epoch in tqdm(range(iterations)): for data in self.dataloader: step = i # inputs data = data.permute(1, 0, 2, 3, 4).to(device) anchor, pos, neg = data # outputs anchor_embedding, spatial_features_anchor = self.model(anchor) pos_embedding, spatial_features_pos = self.model(pos) neg_embedding, spatial_features_neg = self.model(neg) loss = self.contrastive_criterion(anchor_embedding, pos_embedding, neg_embedding) # optimization self.optimizer.zero_grad() loss.backward() self.optimizer.step() # plot loss # normalized_loss=loss.item()/self.contrastive_criterion.margins[self.contrastive_criterion.margin_step] self.writer.add_scalar("training_loss", loss.item(), step) # update margin - step the scheduler if (i + 1) % 1000 == 0 and i > 2000: self.dataset.update_margins() self.contrastive_criterion.update_margin() if (i + 1) % 1000 == 0 or i == 0: for j in range(len(self.images)): img = self.dataset.transform( self.images[j]).unsqueeze(0).to(device) with torch.no_grad(): self.embedding[j], _ = self.model(img) for j in range(len(self.images)): dis = torch.abs(self.embedding[-1] - self.embedding[j]).pow(2).sum(-1) cost = -dis self.writer.add_scalar("cost_%d" % self.cost_counter, cost, j) self.writer.add_embedding(self.embedding, global_step=step) self.cost_counter += 1 i += 1 # soft update if self.load: for param1, param2 in zip(self.model.parameters(), self.delayed_model.parameters()): param2.data.copy_(self.tau * param1.data + (1 - self.tau) * param2.data) param1.data.copy_(param2) # save the model and visualize the embeddings torch.save(self.model.state_dict(), self.save_to) def train_descriptors(self, iterations=150): self.dataset.count = 0 i = 0 for epoch in tqdm(range(iterations)): for data in self.descriptor_dataloader: step = i # inputs data = data.permute(1, 0, 2, 3, 4).to(device) img1, img2, _ = data # outputs embedding1, spatial_features1 = self.model(img1) embedding2, spatial_features2 = self.model(img2) # to be used in pixel loss img1_out = self.descriptor_model(spatial_features1) vis1 = self.dataset.unormalize(img1.squeeze()) vis2 = self.dataset.unormalize(img1_out.squeeze()) img1_un = vis1.permute(1, 2, 0) img2_out = self.descriptor_model(spatial_features2) img2_un = self.dataset.unormalize(img2.squeeze()).permute( 1, 2, 0) if (i + 1) % 1000 == 0: self.writer.add_image( "original vs. prediction", torchvision.utils.make_grid([vis1, vis2]), step) if i > 5000: self.dataset.update_margins() self.pixel_criterion.update_margins() # loss computation loss_p = self.pixel_criterion(img1_un, img1_out, img2_un, img2_out) # optimization self.descriptor_optimizer.zero_grad() loss_p.backward() self.descriptor_optimizer.step() # plot loss # normalized_loss=loss_p.item()/self.pixel_criterion.margins[self.pixel_criterion.margin_step] self.writer.add_scalar("training_descriptor_loss", loss_p.item(), step) torch.save(self.descriptor_model.state_dict(), self.d_save_to) i += 1 # soft update if self.load: for param1, param2 in zip(self.descriptor_model.parameters(), self.d_delayed_model.parameters()): param2.data.copy_(self.tau * param1.data + (1 - self.tau) * param2.data) param1.data.copy_(param2)
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_dataset(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 = False trainloader = DataLoader( VideoDataset(dataset.train, seq_len=args.seq_len, sample='random', 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( VideoDataset(dataset.query, seq_len=args.seq_len, sample='dense', 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='dense', 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: {:.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 = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) if args.stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch start_time = time.time() print(start_time) for batch_idx, (imgs, pids, _) in enumerate(trainloader): print(batch_idx) print('x') if use_gpu: imgs, pids = imgs.cuda(), pids.cuda() imgs, pids = Variable(imgs), Variable(pids) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
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(): runId = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S') cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_DIR, runId) if not os.path.exists(cfg.OUTPUT_DIR): os.mkdir(cfg.OUTPUT_DIR) print(cfg.OUTPUT_DIR) torch.manual_seed(cfg.RANDOM_SEED) random.seed(cfg.RANDOM_SEED) np.random.seed(cfg.RANDOM_SEED) os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID use_gpu = torch.cuda.is_available() and cfg.MODEL.DEVICE == "cuda" if not cfg.EVALUATE_ONLY: sys.stdout = Logger(osp.join(cfg.OUTPUT_DIR, 'log_train.txt')) else: sys.stdout = Logger(osp.join(cfg.OUTPUT_DIR, 'log_test.txt')) print("==========\nConfigs:{}\n==========".format(cfg)) if use_gpu: print("Currently using GPU {}".format(cfg.MODEL.DEVICE_ID)) cudnn.benchmark = True torch.cuda.manual_seed_all(cfg.RANDOM_SEED) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(cfg.DATASETS.NAME)) dataset = data_manager.init_dataset(root=cfg.DATASETS.ROOT_DIR, name=cfg.DATASETS.NAME) print("Initializing model: {}".format(cfg.MODEL.NAME)) if cfg.MODEL.ARCH == 'video_baseline': torch.backends.cudnn.benchmark = False model = models.init_model(name=cfg.MODEL.ARCH, num_classes=625, pretrain_choice=cfg.MODEL.PRETRAIN_CHOICE, last_stride=cfg.MODEL.LAST_STRIDE, neck=cfg.MODEL.NECK, model_name=cfg.MODEL.NAME, neck_feat=cfg.TEST.NECK_FEAT, model_path=cfg.MODEL.PRETRAIN_PATH) print("Model size: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) transform_train = T.Compose([ T.Resize(cfg.INPUT.SIZE_TRAIN), T.RandomHorizontalFlip(p=cfg.INPUT.PROB), T.Pad(cfg.INPUT.PADDING), T.RandomCrop(cfg.INPUT.SIZE_TRAIN), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), T.RandomErasing(probability=cfg.INPUT.RE_PROB, mean=cfg.INPUT.PIXEL_MEAN) ]) transform_test = T.Compose([ T.Resize(cfg.INPUT.SIZE_TEST), 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 cfg.DATALOADER.NUM_WORKERS = 0 trainloader = DataLoader(VideoDataset( dataset.train, seq_len=cfg.DATASETS.SEQ_LEN, sample=cfg.DATASETS.TRAIN_SAMPLE_METHOD, transform=transform_train, dataset_name=cfg.DATASETS.NAME), sampler=RandomIdentitySampler( dataset.train, num_instances=cfg.DATALOADER.NUM_INSTANCE), batch_size=cfg.SOLVER.SEQS_PER_BATCH, num_workers=cfg.DATALOADER.NUM_WORKERS, pin_memory=pin_memory, drop_last=True) queryloader = DataLoader(VideoDataset( dataset.query, seq_len=cfg.DATASETS.SEQ_LEN, sample=cfg.DATASETS.TEST_SAMPLE_METHOD, transform=transform_test, max_seq_len=cfg.DATASETS.TEST_MAX_SEQ_NUM, dataset_name=cfg.DATASETS.NAME), batch_size=cfg.TEST.SEQS_PER_BATCH, shuffle=False, num_workers=cfg.DATALOADER.NUM_WORKERS, pin_memory=pin_memory, drop_last=False) galleryloader = DataLoader( VideoDataset(dataset.gallery, seq_len=cfg.DATASETS.SEQ_LEN, sample=cfg.DATASETS.TEST_SAMPLE_METHOD, transform=transform_test, max_seq_len=cfg.DATASETS.TEST_MAX_SEQ_NUM, dataset_name=cfg.DATASETS.NAME), batch_size=cfg.TEST.SEQS_PER_BATCH, shuffle=False, num_workers=cfg.DATALOADER.NUM_WORKERS, pin_memory=pin_memory, drop_last=False, ) if cfg.MODEL.SYN_BN: if use_gpu: model = nn.DataParallel(model) if cfg.SOLVER.FP_16: model = apex.parallel.convert_syncbn_model(model) model.cuda() start_time = time.time() xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids) tent = TripletLoss(cfg.SOLVER.MARGIN) optimizer = make_optimizer(cfg, model) scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS, cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR, cfg.SOLVER.WARMUP_ITERS, cfg.SOLVER.WARMUP_METHOD) # metrics = test(model, queryloader, galleryloader, cfg.TEST.TEMPORAL_POOL_METHOD, use_gpu) no_rise = 0 best_rank1 = 0 start_epoch = 0 for epoch in range(start_epoch, cfg.SOLVER.MAX_EPOCHS): # if no_rise == 10: # break scheduler.step() print("noriase:", no_rise) print("==> Epoch {}/{}".format(epoch + 1, cfg.SOLVER.MAX_EPOCHS)) print("current lr:", scheduler.get_lr()[0]) train(model, trainloader, xent, tent, optimizer, use_gpu) if cfg.SOLVER.EVAL_PERIOD > 0 and ( (epoch + 1) % cfg.SOLVER.EVAL_PERIOD == 0 or (epoch + 1) == cfg.SOLVER.MAX_EPOCHS): print("==> Test") metrics = test(model, queryloader, galleryloader, cfg.TEST.TEMPORAL_POOL_METHOD, use_gpu) rank1 = metrics[0] if rank1 > best_rank1: best_rank1 = rank1 no_rise = 0 else: no_rise += 1 continue if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() torch.save( state_dict, osp.join( cfg.OUTPUT_DIR, "rank1_" + str(rank1) + '_checkpoint_ep' + str(epoch + 1) + '.pth')) # best_p = osp.join(cfg.OUTPUT_DIR, "rank1_" + str(rank1) + '_checkpoint_ep' + str(epoch + 1) + '.pth') elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
# strategy = tf.distribute.MirroredStrategy() # print("Number of devices: {}".format(strategy.num_replicas_in_sync)) # TODO add lambda regularization # TODO other regularizer (batch norm, dropout etc) # TODO add activation layers ? print(f"image data format is {tf.keras.backend.image_data_format()}") latest_ckpt = tf.train.latest_checkpoint(CHECKPOINT_DIR) datasets = data.get_hard_images() model = DeepRankingNetwork(IMG_SHAPE, EMBEDDINGS_DIM) base_lr = 0.001 max_lr = 0.06 model.compile( optimizer=tf.keras.optimizers.Adam(base_lr), # TODO choose optimizer loss=TripletLoss(GAP_PARAMETER, EMBEDDINGS_DIM, BATCH_SIZE), metrics=[], ) if latest_ckpt is not None: model.load_weights(latest_ckpt) model.fit( datasets["train"][1], validation_data=datasets["validation"][1], validation_steps=datasets["validation"][0] // BATCH_SIZE, steps_per_epoch=datasets["train"][0] // BATCH_SIZE, epochs=EPOCHS, callbacks=[ tf.keras.callbacks.LearningRateScheduler( schedule=lambda epoch_index: base_lr + (max_lr - base_lr) *
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_dataset(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_train_p = T.Compose([ T.Random2DTranslation(256, 128), 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]), ]) transform_test_p = T.Compose([ T.Resize((256, 128)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) train_file = 'data/cuhk_train.pkl' test_file = 'data/cuhk_test.pkl' gallery_file = 'data/cuhk_gallery.pkl' data_root = args.data_root dataset_train = CUHKGroup(train_file, data_root, True, transform_train, transform_train_p) dataset_test = CUHKGroup(test_file, data_root, False, transform_test, transform_test_p) dataset_query = CUHKGroup(test_file, data_root, False, transform_test, transform_test_p) dataset_gallery = CUHKGroup(gallery_file, data_root, False, transform_test, transform_test_p) pin_memory = True if use_gpu else False if args.xent_only: trainloader = DataLoader( dataset_train, batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) else: trainloader = DataLoader( dataset_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( dataset_test, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) querygalleryloader = DataLoader( dataset_query, batch_size=args.gallery_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) galleryloader = DataLoader( dataset_gallery, batch_size=args.gallery_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) print("Initializing model: {}".format(args.arch)) if args.xent_only: # model = models.init_model(name=args.arch, num_classes=dataset_train.num_train_gids, loss={'xent'}) model = models.init_model(name=args.arch, num_classes=dataset_train.num_train_gids, loss={'xent'}) else: # model = models.init_model(name=args.arch, num_classes=dataset_train.num_train_gids, loss={'xent', 'htri'}) model = models.init_model( name=args.arch, num_classes=dataset_train.num_train_gids, num_person_classes=dataset_train.num_train_pids, loss={'xent', 'htri'}) #criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset_train.num_train_gids, use_gpu=use_gpu) #criterion_xent_person = CrossEntropyLabelSmooth(num_classes=dataset_train.num_train_pids, use_gpu=use_gpu) if os.path.exists(args.pretrained_model): print("Loading checkpoint from '{}'".format(args.pretrained_model)) checkpoint = torch.load(args.pretrained_model) model_dict = model.state_dict() pretrain_dict = checkpoint['state_dict'] pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) criterion_xent = nn.CrossEntropyLoss(ignore_index=-1) criterion_xent_person = nn.CrossEntropyLoss(ignore_index=-1) criterion_htri = TripletLoss(margin=args.margin) criterion_pair = ContrastiveLoss(margin=args.margin) criterion_htri_filter = TripletLossFilter(margin=args.margin) criterion_permutation = PermutationLoss() optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) if args.stepsize > 0: if args.warmup: scheduler = WarmupMultiStepLR(optimizer, [200, 400, 600]) else: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test_gcn_person_batch(model, queryloader, querygalleryloader, galleryloader, args.pool, use_gpu) #test_gcn_batch(model, queryloader, querygalleryloader, galleryloader, args.pool, use_gpu) #test_gcn(model, queryloader, galleryloader, args.pool, use_gpu) #test(model, queryloader, galleryloader, args.pool, use_gpu) return start_time = time.time() best_rank1 = -np.inf for epoch in range(start_epoch, args.max_epoch): #print("==> Epoch {}/{} lr:{}".format(epoch + 1, args.max_epoch, scheduler.get_lr()[0])) train_gcn(model, criterion_xent, criterion_xent_person, criterion_pair, criterion_htri_filter, criterion_htri, criterion_permutation, optimizer, trainloader, use_gpu) #train(model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) if args.stepsize > 0: scheduler.step() if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or ( epoch + 1) == args.max_epoch: print("==> Test") rank1 = test_gcn_person_batch(model, queryloader, querygalleryloader, galleryloader, args.pool, use_gpu) #rank1 = test_gcn(model, queryloader, galleryloader, args.pool, use_gpu=False) #rank1 = test(model, queryloader, galleryloader, args.pool, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 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')) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
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_dataset(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 trainloader = DataLoader( VideoDataset(dataset.train, seq_len=args.seq_len, sample='random', 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( VideoDataset(dataset.query, seq_len=args.seq_len, sample='dense', 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='dense', 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)) if args.arch == 'resnet503d': model = resnet3d.resnet50(num_classes=dataset.num_train_pids, sample_width=args.width, sample_height=args.height, sample_duration=args.seq_len) if not os.path.exists(args.pretrained_model): raise IOError("Can't find pretrained model: {}".format( args.pretrained_model)) print("Loading checkpoint from '{}'".format(args.pretrained_model)) checkpoint = torch.load(args.pretrained_model) state_dict = {} for key in checkpoint['state_dict']: if 'fc' in key: continue state_dict[key.partition("module.") [2]] = checkpoint['state_dict'][key] model.load_state_dict(state_dict, strict=False) else: model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) 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 = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) if args.stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) start_epoch = args.start_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() best_rank1 = -np.inf if args.arch == 'resnet503d': torch.backends.cudnn.benchmark = False for epoch in range(start_epoch, args.max_epoch): print("==> Epoch {}/{}".format(epoch + 1, args.max_epoch)) train(model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) if args.stepsize > 0: scheduler.step() if 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 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')) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
# {"criterion": CrossEntropyLabelSmooth(num_classes=class_num), "weight": 1.0,}, # {"criterion": TripletLoss(margin="0.3", metric="euclidean"), "weight": 1.0,}, { "criterion": TripletLoss_WRT(), "weight": 1.0, }, ]) else: criterion = Criterion([ { "criterion": nn.CrossEntropyLoss(), "weight": 1.0, }, # {"criterion": CrossEntropyLabelSmooth(num_classes=class_num), "weight": 1.0,}, { "criterion": TripletLoss(margin="0.3", metric="euclidean"), "weight": 1.0, }, ]) # build optimizer if args.optim == "sgd": args.lr = 0.1 ignored_params = list(map(id, net.bottleneck.parameters())) + list( map(id, net.classifier.parameters())) base_params = filter(lambda p: id(p) not in ignored_params, net.parameters()) optimizer = torch.optim.SGD( [ { "params": base_params,
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 # set a learning rate if args.lr_factor == -1: args.lr_factor = random() args.lr = args.lr_factor * 10**-args.lr_base #print(f"Choose learning rate {args.lr}") sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'), mode='a') print("==========\nArgs:{}\n==========".format(args)) #assert torch.distributed.is_available() #print("Initializing DDP by nccl-tcp({}) rank({}) world_size({})".format(args.init_method, args.rank, args.world_size)) #dist.init_process_group(backend='nccl', init_method=args.init_method, rank=args.rank, world_size=args.world_size) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) 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_dataset(name=args.dataset, root=args.root) # Data augmentation spatial_transform_train = [ ST.Scale((args.height, args.width), interpolation=3), ST.RandomHorizontalFlip(), ST.ToTensor(), ST.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] spatial_transform_train = ST.Compose(spatial_transform_train) temporal_transform_train = TT.TemporalRandomCrop(size=args.seq_len, stride=args.sample_stride) #temporal_transform_train = TT.TemporalRandomCropPick(size=args.seq_len, stride=args.sample_stride) spatial_transform_test = ST.Compose([ ST.Scale((args.height, args.width), interpolation=3), ST.ToTensor(), ST.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) temporal_transform_test = TT.TemporalBeginCrop(size=args.test_frames) pin_memory = True if use_gpu else False dataset_train = dataset.train if args.dataset == 'duke': dataset_train = dataset.train_dense print('process duke dataset') #sampler = RandomIdentitySampler(dataset_train, num_instances=args.num_instances) if args.dataset == 'lsvid': sampler = RandomIdentitySampler(dataset_train, num_instances=args.num_instances) elif args.dataset == 'mars': sampler = RandomIdentitySampler(dataset_train, num_instances=args.num_instances) trainloader = DataLoader( VideoDataset(dataset_train, spatial_transform=spatial_transform_train, temporal_transform=temporal_transform_train), sampler=sampler, batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) ''' for batch_idx, (vids, pids, camids, img_paths) in enumerate(trainloader): print(batch_idx, pids, camids, img_paths) break return ''' dataset_query = dataset.query dataset_gallery = dataset.gallery if args.dataset == 'lsvid': dataset_query = dataset.val_query dataset_gallery = dataset.val_gallery print('process lsvid dataset') queryloader = DataLoader( VideoDataset(dataset_query, spatial_transform=spatial_transform_test, temporal_transform=temporal_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, spatial_transform=spatial_transform_test, temporal_transform=temporal_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, use_gpu=use_gpu, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) #print(model) if args.resume: print("Loading checkpoint from '{}'".format(args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) criterion_xent = nn.CrossEntropyLoss() criterion_htri = TripletLoss(margin=args.margin, distance=args.distance, use_gpu=use_gpu) criterion_htri_c = TripletInterCamLoss(margin=args.margin, distance=args.distance, use_gpu=use_gpu) #criterion_htri_c = TripletWeightedInterCamLoss(margin=args.margin, distance=args.distance, use_gpu=use_gpu, alpha=args.cam_alpha) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if use_gpu: model = nn.DataParallel(model).cuda() #model = model.cuda() #model = nn.parallel.DistributedDataParallel(model) 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): #print("Set sampler seed to {}".format(args.seed*epoch)) #sampler.set_seed(args.seed*epoch) start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, criterion_htri_c, 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") with torch.no_grad(): 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 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))
print('datasets loaded') train_loader = torch.utils.data.DataLoader(train_set, batch_size=cfg.data.imgs_per_gpu, shuffle=True) query_loader = torch.utils.data.DataLoader(query_set, batch_size=cfg.data.imgs_per_gpu, shuffle=True) print('dataloader created') # Build model backbone_model = Vgg16L2(num_dim=128) model = TripletNet(backbone_model) model.cuda() print('model built') margin = 1. loss_fn = TripletLoss(margin) lr = 1e-4 optimizer = optim.Adam(model.parameters(), lr=lr) scheduler = lr_scheduler.StepLR(optimizer, 8, gamma=0.1, last_epoch=-1) n_epochs = 20 log_interval = 500 print('start training') fit(train_loader, query_loader, model, loss_fn, optimizer, scheduler, n_epochs, cuda, log_interval) print('done training') torch.save(model.state_dict(), './checkpoint/triplet_vanilla_20_epochs_1e-4.pth')
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 sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'), mode='a') print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) 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_dataset(name=args.dataset, root=args.root) # Data augmentation spatial_transform_train = [ ST.Scale((args.height, args.width), interpolation=3), ST.RandomHorizontalFlip(), ST.ToTensor(), ST.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ] spatial_transform_train = ST.Compose(spatial_transform_train) temporal_transform_train = TT.TemporalRandomCrop(size=args.seq_len, stride=args.sample_stride) spatial_transform_test = ST.Compose([ ST.Scale((args.height, args.width), interpolation=3), ST.ToTensor(), ST.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) temporal_transform_test = TT.TemporalBeginCrop(size=args.test_frames) pin_memory = True if use_gpu else False dataset_train = dataset.train if args.dataset == 'duke': dataset_train = dataset.train_dense print('process duke dataset') trainloader = DataLoader( VideoDataset(dataset_train, spatial_transform=spatial_transform_train, temporal_transform=temporal_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(VideoDataset( dataset.query, spatial_transform=spatial_transform_test, temporal_transform=temporal_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, spatial_transform=spatial_transform_test, temporal_transform=temporal_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, use_gpu=use_gpu, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) print(model) criterion_xent = nn.CrossEntropyLoss() criterion_htri = TripletLoss(margin=args.margin, distance=args.distance, use_gpu=use_gpu) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if use_gpu: model = nn.DataParallel(model).cuda() 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 ( epoch + 1) % args.eval_step == 0 or epoch == 0: print("==> Test") with torch.no_grad(): 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))