def get_extractor(): cfg = Config() TVT, TMO = set_devices(cfg.sys_device_ids) if cfg.seed is not None: set_seed(cfg.seed) ####################### # print configuration # ####################### import pprint print('-' * 60) print('cfg.__dict__') pprint.pprint(cfg.__dict__) print('-' * 60) ########## # Models # ########## model = Model() # num_class is None, no need for local feature model_w = DataParallel(model) print(model.state_dict().keys()) # Transfer model to specify devices TMO([model]) assert (cfg.model_weight_file != '') map_location = (lambda storage, loc: storage) sd = torch.load(cfg.model_weight_file, map_location=map_location) if 'state_dicts' in sd: # load from checkpoint print(len(sd['state_dicts'])) sd = sd['state_dicts'][0] load_state_dict(model, sd) print('Loaded model weights from {}'.format(cfg.model_weight_file)) ########### # Extract # ########### extractor = ExtractFeature(model_w, TVT) return extractor
def __init__(self, ): os.chdir('./src/person_reid/AlignedReID-Re-Production-Pytorch-master/') import torch from torch.nn.parallel import DataParallel from aligned_reid.model.Model import Model from aligned_reid.utils.utils import load_state_dict from aligned_reid.utils.utils import set_devices os.chdir('../../../') model = Model(local_conv_out_channels=128, num_classes=751) # Model wrapper model_w = DataParallel(model) weight = './model/align_reid/model_weight.pth' map_location = (lambda storage, loc: storage) sd = torch.load(weight, map_location=map_location) load_state_dict(model, sd) print('Loaded model weights from {}'.format(weight)) TVT, TMO = set_devices((0, )) self.featureExtractor = AlignReid_Embedder.ExtractFeature(model_w, TVT)
def main(): # name_data = '20180525184611494' name_data = '20180717103311494' TVT, TMO = set_devices((0,)) dataset_kwargs = dict( name = name_data, resize_h_w = (256, 128), scale = True, im_mean = [0.486, 0.459, 0.408], im_std = [0.229, 0.224, 0.225], batch_dims = 'NCHW', num_prefetch_threads = 1) test_set_kwargs = dict( part = 'val', batch_size = 32, final_batch = True, shuffle = False, mirror_type = ['random', 'always', None][2], prng = np.random) test_set_kwargs.update(dataset_kwargs) test_set = create_dataset(**test_set_kwargs) with measure_time('Load model'): model = Model(local_conv_out_channels = 128, num_classes = 1000) model_w = DataParallel(model) optimizer = optim.Adam(model.parameters(), lr = 2e-4, weight_decay = 0.0005) modules_optims = [model, optimizer] with measure_time('Load checkpoint'): load_ckpt(modules_optims, 'ckpt.pth') test_set.set_feat_func(ExtractFeature(model_w, TVT)) exp_summary(test_set)
def main(): cfg = Config() # Redirect logs to both console and file. if cfg.log_to_file: ReDirectSTD(cfg.stdout_file, 'stdout', False) ReDirectSTD(cfg.stderr_file, 'stderr', False) # Lazily create SummaryWriter writer = None TVT, TMO = set_devices(cfg.sys_device_ids) if cfg.seed is not None: set_seed(cfg.seed) # Dump the configurations to log. import pprint print('-' * 60) print('cfg.__dict__') pprint.pprint(cfg.__dict__) print('-' * 60) ########### # Dataset # ########### train_set = create_dataset(**cfg.train_set_kwargs) test_sets = [] test_set_names = [] if cfg.dataset == 'combined': for name in ['market1501', 'cuhk03', 'duke']: cfg.test_set_kwargs['name'] = name test_sets.append(create_dataset(**cfg.test_set_kwargs)) test_set_names.append(name) else: test_sets.append(create_dataset(**cfg.test_set_kwargs)) test_set_names.append(cfg.dataset) ########### # Models # ########### model = Model(local_conv_out_channels=cfg.local_conv_out_channels, num_classes=len(train_set.ids2labels)) # Model wrapper model_w = DataParallel(model) ############################# # Criteria and Optimizers # ############################# id_criterion = nn.CrossEntropyLoss() g_tri_loss = TripletLoss(margin=cfg.global_margin) l_tri_loss = TripletLoss(margin=cfg.local_margin) optimizer = optim.Adam(model.parameters(), lr=cfg.base_lr, weight_decay=cfg.weight_decay) # Bind them together just to save some codes in the following usage. modules_optims = [model, optimizer] ################################ # May Resume Models and Optims # ################################ if cfg.resume: resume_ep, scores = load_ckpt(modules_optims, cfg.ckpt_file) # May Transfer Models and Optims to Specified Device. Transferring optimizer # is to cope with the case when you load the checkpoint to a new device. TMO(modules_optims) ######## # Test # ######## def test(load_model_weight=False): if load_model_weight: if cfg.model_weight_file != '': map_location = (lambda storage, loc: storage) sd = torch.load(cfg.model_weight_file, map_location=map_location) load_state_dict(model, sd) print('Loaded model weights from {}'.format(cfg.model_weight_file)) else: load_ckpt(modules_optims, cfg.ckpt_file) use_local_distance = (cfg.l_loss_weight > 0) \ and cfg.local_dist_own_hard_sample for test_set, name in zip(test_sets, test_set_names): test_set.set_feat_func(ExtractFeature(model_w, TVT)) print('\n=========> Test on dataset: {} <=========\n'.format(name)) test_set.eval( normalize_feat=cfg.normalize_feature, use_local_distance=use_local_distance) if cfg.only_test: test(load_model_weight=True) return ############ # Training # ############ start_ep = resume_ep if cfg.resume else 0 for ep in range(start_ep, cfg.total_epochs): # Adjust Learning Rate if cfg.lr_decay_type == 'exp': adjust_lr_exp( optimizer, cfg.base_lr, ep + 1, cfg.total_epochs, cfg.exp_decay_at_epoch) else: adjust_lr_staircase( optimizer, cfg.base_lr, ep + 1, cfg.staircase_decay_at_epochs, cfg.staircase_decay_multiply_factor) may_set_mode(modules_optims, 'train') g_prec_meter = AverageMeter() g_m_meter = AverageMeter() g_dist_ap_meter = AverageMeter() g_dist_an_meter = AverageMeter() g_loss_meter = AverageMeter() l_prec_meter = AverageMeter() l_m_meter = AverageMeter() l_dist_ap_meter = AverageMeter() l_dist_an_meter = AverageMeter() l_loss_meter = AverageMeter() id_loss_meter = AverageMeter() loss_meter = AverageMeter() ep_st = time.time() step = 0 epoch_done = False while not epoch_done: step += 1 step_st = time.time() ims, im_names, labels, mirrored, epoch_done = train_set.next_batch() ims_var = Variable(TVT(torch.from_numpy(ims).float())) labels_t = TVT(torch.from_numpy(labels).long()) labels_var = Variable(labels_t) global_feat, local_feat, logits = model_w(ims_var) g_loss, p_inds, n_inds, g_dist_ap, g_dist_an, g_dist_mat = global_loss( g_tri_loss, global_feat, labels_t, normalize_feature=cfg.normalize_feature) if cfg.l_loss_weight == 0: l_loss = 0 elif cfg.local_dist_own_hard_sample: # Let local distance find its own hard samples. l_loss, l_dist_ap, l_dist_an, _ = local_loss( l_tri_loss, local_feat, None, None, labels_t, normalize_feature=cfg.normalize_feature) else: l_loss, l_dist_ap, l_dist_an = local_loss( l_tri_loss, local_feat, p_inds, n_inds, labels_t, normalize_feature=cfg.normalize_feature) id_loss = 0 if cfg.id_loss_weight > 0: id_loss = id_criterion(logits, labels_var) loss = g_loss * cfg.g_loss_weight \ + l_loss * cfg.l_loss_weight \ + id_loss * cfg.id_loss_weight optimizer.zero_grad() loss.backward() optimizer.step() ############ # Step Log # ############ # precision g_prec = (g_dist_an > g_dist_ap).data.float().mean() # the proportion of triplets that satisfy margin g_m = (g_dist_an > g_dist_ap + cfg.global_margin).data.float().mean() g_d_ap = g_dist_ap.data.mean() g_d_an = g_dist_an.data.mean() g_prec_meter.update(g_prec) g_m_meter.update(g_m) g_dist_ap_meter.update(g_d_ap) g_dist_an_meter.update(g_d_an) g_loss_meter.update(to_scalar(g_loss)) if cfg.l_loss_weight > 0: # precision l_prec = (l_dist_an > l_dist_ap).data.float().mean() # the proportion of triplets that satisfy margin l_m = (l_dist_an > l_dist_ap + cfg.local_margin).data.float().mean() l_d_ap = l_dist_ap.data.mean() l_d_an = l_dist_an.data.mean() l_prec_meter.update(l_prec) l_m_meter.update(l_m) l_dist_ap_meter.update(l_d_ap) l_dist_an_meter.update(l_d_an) l_loss_meter.update(to_scalar(l_loss)) if cfg.id_loss_weight > 0: id_loss_meter.update(to_scalar(id_loss)) loss_meter.update(to_scalar(loss)) if step % cfg.log_steps == 0: time_log = '\tStep {}/Ep {}, {:.2f}s'.format( step, ep + 1, time.time() - step_st, ) if cfg.g_loss_weight > 0: g_log = (', gp {:.2%}, gm {:.2%}, ' 'gd_ap {:.4f}, gd_an {:.4f}, ' 'gL {:.4f}'.format( g_prec_meter.val, g_m_meter.val, g_dist_ap_meter.val, g_dist_an_meter.val, g_loss_meter.val, )) else: g_log = '' if cfg.l_loss_weight > 0: l_log = (', lp {:.2%}, lm {:.2%}, ' 'ld_ap {:.4f}, ld_an {:.4f}, ' 'lL {:.4f}'.format( l_prec_meter.val, l_m_meter.val, l_dist_ap_meter.val, l_dist_an_meter.val, l_loss_meter.val, )) else: l_log = '' if cfg.id_loss_weight > 0: id_log = (', idL {:.4f}'.format(id_loss_meter.val)) else: id_log = '' total_loss_log = ', loss {:.4f}'.format(loss_meter.val) log = time_log + \ g_log + l_log + id_log + \ total_loss_log print(log) ############# # Epoch Log # ############# time_log = 'Ep {}, {:.2f}s'.format(ep + 1, time.time() - ep_st, ) if cfg.g_loss_weight > 0: g_log = (', gp {:.2%}, gm {:.2%}, ' 'gd_ap {:.4f}, gd_an {:.4f}, ' 'gL {:.4f}'.format( g_prec_meter.avg, g_m_meter.avg, g_dist_ap_meter.avg, g_dist_an_meter.avg, g_loss_meter.avg, )) else: g_log = '' if cfg.l_loss_weight > 0: l_log = (', lp {:.2%}, lm {:.2%}, ' 'ld_ap {:.4f}, ld_an {:.4f}, ' 'lL {:.4f}'.format( l_prec_meter.avg, l_m_meter.avg, l_dist_ap_meter.avg, l_dist_an_meter.avg, l_loss_meter.avg, )) else: l_log = '' if cfg.id_loss_weight > 0: id_log = (', idL {:.4f}'.format(id_loss_meter.avg)) else: id_log = '' total_loss_log = ', loss {:.4f}'.format(loss_meter.avg) log = time_log + \ g_log + l_log + id_log + \ total_loss_log print(log) # Log to TensorBoard if cfg.log_to_file: if writer is None: writer = SummaryWriter(log_dir=osp.join(cfg.exp_dir, 'tensorboard')) writer.add_scalars( 'loss', dict(global_loss=g_loss_meter.avg, local_loss=l_loss_meter.avg, id_loss=id_loss_meter.avg, loss=loss_meter.avg, ), ep) writer.add_scalars( 'tri_precision', dict(global_precision=g_prec_meter.avg, local_precision=l_prec_meter.avg, ), ep) writer.add_scalars( 'satisfy_margin', dict(global_satisfy_margin=g_m_meter.avg, local_satisfy_margin=l_m_meter.avg, ), ep) writer.add_scalars( 'global_dist', dict(global_dist_ap=g_dist_ap_meter.avg, global_dist_an=g_dist_an_meter.avg, ), ep) writer.add_scalars( 'local_dist', dict(local_dist_ap=l_dist_ap_meter.avg, local_dist_an=l_dist_an_meter.avg, ), ep) # save ckpt if cfg.log_to_file: save_ckpt(modules_optims, ep + 1, 0, cfg.ckpt_file) ######## # Test # ######## test(load_model_weight=False)
def __call__(self, ims): print(len(ims)) old_train_eval_model = self.model.training # Set eval mode. # Force all BN layers to use global mean and variance, also disable # dropout. self.model.eval() ims = Variable(self.TVT(torch.from_numpy(ims).float())) global_feat, local_feat = self.model(ims)[:2] global_feat = global_feat.data.cpu().numpy() local_feat = local_feat.data.cpu().numpy() # Restore the model to its old train/eval mode. self.model.train(old_train_eval_model) return global_feat, local_feat model = Model(local_conv_out_channels=128, num_classes=751) # Model wrapper model_w = DataParallel(model) weight = './model_weight.pth' map_location = (lambda storage, loc: storage) sd = torch.load(weight, map_location=map_location) load_state_dict(model, sd) print('Loaded model weights from {}'.format(weight)) TVT, TMO = set_devices((0, )) FeatureExtractor = ExtractFeature(model_w, TVT) global_feat, local_feat = FeatureExtractor(XXX)
weight_decay = 0.0005 optimizer = optim.Adam(model.parameters(), lr=base_lr, weight_decay=weight_decay) # Bind them together just to save some codes in the following usage. modules_optims = [model, optimizer] model_weight_file = '../../model_weight.pth' map_location = (lambda storage, loc: storage) sd = torch.load(model_weight_file, map_location=map_location) load_state_dict(model, sd) print('Loaded model weights from {}'.format(model_weight_file)) sys_device_ids = (0,) TVT, TMO = set_devices(sys_device_ids) name_list = ['gow.jpg', 'gunes.jpg', 'niru.jpg', 'gow_query.jpg'] global_features_list = [] local_features_list = [] for name in name_list: input_image = cv2.imread('my_images/' + name) resized_image = cv2.resize(input_image, (128, 256)) # input_image = np.asarray(Image.open('gow_query.jpg')) transposed = resized_image.transpose(2,0,1) test_img = transposed[np.newaxis]
def main(): # reranking_mAP_list = [] cfg = Config() # Redirect logs to both console and file. if cfg.log_to_file: ReDirectSTD(cfg.stdout_file, 'stdout', False) ReDirectSTD(cfg.stderr_file, 'stderr', False) # Lazily create SummaryWriter writer = None TVT, TMO = set_devices(cfg.sys_device_ids) if cfg.seed is not None: set_seed(cfg.seed) # Dump the configurations to log. import pprint print('-' * 60) print('cfg.__dict__') pprint.pprint(cfg.__dict__) print('-' * 60) ########### # Dataset # ########### train_set = create_dataset(**cfg.train_set_kwargs) test_sets = [] test_set_names = [] if cfg.dataset == 'combined': for name in ['market1501', 'cuhk03', 'duke']: cfg.test_set_kwargs['name'] = name test_sets.append(create_dataset(**cfg.test_set_kwargs)) test_set_names.append(name) else: test_sets.append(create_dataset(**cfg.test_set_kwargs)) test_set_names.append(cfg.dataset) ########### # Models # ########### model = Model(local_conv_out_channels=cfg.local_conv_out_channels, num_classes=len(train_set.ids2labels)) # Model wrapper model_w = DataParallel(model) ############################# # Criteria and Optimizers # ############################# id_criterion = nn.CrossEntropyLoss() g_tri_loss = TripletLoss(margin=cfg.global_margin) l_tri_loss = TripletLoss(margin=cfg.local_margin) optimizer = optim.Adam(model.parameters(), lr=cfg.base_lr, weight_decay=cfg.weight_decay) # Bind them together just to save some codes in the following usage. modules_optims = [model, optimizer] ################################ # May Resume Models and Optims # ################################ if cfg.resume: resume_ep, scores = load_ckpt(modules_optims, cfg.ckpt_file) # May Transfer Models and Optims to Specified Device. Transferring optimizer # is to cope with the case when you load the checkpoint to a new device. TMO(modules_optims) ######## # Test # ######## def test(load_model_weight=False): if load_model_weight: if cfg.model_weight_file != '': map_location = (lambda storage, loc: storage) sd = torch.load(cfg.model_weight_file, map_location=map_location) load_state_dict(model, sd) print('Loaded model weights from {}'.format( cfg.model_weight_file)) else: load_ckpt(modules_optims, cfg.ckpt_file) use_local_distance = (cfg.l_loss_weight > 0) \ and cfg.local_dist_own_hard_sample for test_set, name in zip(test_sets, test_set_names): test_set.set_feat_func(ExtractFeature(model_w, TVT)) print('\n=========> Test on dataset: {} <=========\n'.format(name)) test_set.eval(normalize_feat=cfg.normalize_feature, use_local_distance=use_local_distance) # reranking_mAP_list.append(mAP) if cfg.only_test: test(load_model_weight=True) return ############ # Training # ############ start_ep = resume_ep if cfg.resume else 0 for ep in range(start_ep, cfg.total_epochs): # Adjust Learning Rate if cfg.lr_decay_type == 'exp': adjust_lr_exp(optimizer, cfg.base_lr, ep + 1, cfg.total_epochs, cfg.exp_decay_at_epoch) else: adjust_lr_staircase(optimizer, cfg.base_lr, ep + 1, cfg.staircase_decay_at_epochs, cfg.staircase_decay_multiply_factor) may_set_mode(modules_optims, 'train') g_prec_meter = AverageMeter() g_m_meter = AverageMeter() g_dist_ap_meter = AverageMeter() g_dist_an_meter = AverageMeter() g_loss_meter = AverageMeter() l_prec_meter = AverageMeter() l_m_meter = AverageMeter() l_dist_ap_meter = AverageMeter() l_dist_an_meter = AverageMeter() l_loss_meter = AverageMeter() id_loss_meter = AverageMeter() loss_meter = AverageMeter() ep_st = time.time() step = 0 epoch_done = False while not epoch_done: step += 1 step_st = time.time() ims, im_names, labels, mirrored, epoch_done = train_set.next_batch( ) ims_var = Variable(TVT(torch.from_numpy(ims).float())) labels_t = TVT(torch.from_numpy(labels).long()) labels_var = Variable(labels_t) global_feat, local_feat, logits = model_w(ims_var) g_loss, p_inds, n_inds, g_dist_ap, g_dist_an, g_dist_mat = global_loss( g_tri_loss, global_feat, labels_t, normalize_feature=cfg.normalize_feature) if cfg.l_loss_weight == 0: l_loss = 0 elif cfg.local_dist_own_hard_sample: # Let local distance find its own hard samples. l_loss, l_dist_ap, l_dist_an, _ = local_loss( l_tri_loss, local_feat, None, None, labels_t, normalize_feature=cfg.normalize_feature) else: l_loss, l_dist_ap, l_dist_an = local_loss( l_tri_loss, local_feat, p_inds, n_inds, labels_t, normalize_feature=cfg.normalize_feature) id_loss = 0 if cfg.id_loss_weight > 0: id_loss = id_criterion(logits, labels_var) loss = g_loss * cfg.g_loss_weight \ + l_loss * cfg.l_loss_weight \ + id_loss * cfg.id_loss_weight optimizer.zero_grad() loss.backward() optimizer.step() ############ # Step Log # ############ # precision g_prec = (g_dist_an > g_dist_ap).data.float().mean() # the proportion of triplets that satisfy margin g_m = (g_dist_an > g_dist_ap + cfg.global_margin).data.float().mean() g_d_ap = g_dist_ap.data.mean() g_d_an = g_dist_an.data.mean() g_prec_meter.update(g_prec) g_m_meter.update(g_m) g_dist_ap_meter.update(g_d_ap) g_dist_an_meter.update(g_d_an) g_loss_meter.update(to_scalar(g_loss)) if cfg.l_loss_weight > 0: # precision l_prec = (l_dist_an > l_dist_ap).data.float().mean() # the proportion of triplets that satisfy margin l_m = (l_dist_an > l_dist_ap + cfg.local_margin).data.float().mean() l_d_ap = l_dist_ap.data.mean() l_d_an = l_dist_an.data.mean() l_prec_meter.update(l_prec) l_m_meter.update(l_m) l_dist_ap_meter.update(l_d_ap) l_dist_an_meter.update(l_d_an) l_loss_meter.update(to_scalar(l_loss)) if cfg.id_loss_weight > 0: id_loss_meter.update(to_scalar(id_loss)) loss_meter.update(to_scalar(loss)) if step % cfg.log_steps == 0: time_log = '\tStep {}/Ep {}, {:.2f}s'.format( step, ep + 1, time.time() - step_st, ) if cfg.g_loss_weight > 0: g_log = (', gp {:.2%}, gm {:.2%}, ' 'gd_ap {:.4f}, gd_an {:.4f}, ' 'gL {:.4f}'.format( g_prec_meter.val, g_m_meter.val, g_dist_ap_meter.val, g_dist_an_meter.val, g_loss_meter.val, )) else: g_log = '' if cfg.l_loss_weight > 0: l_log = (', lp {:.2%}, lm {:.2%}, ' 'ld_ap {:.4f}, ld_an {:.4f}, ' 'lL {:.4f}'.format( l_prec_meter.val, l_m_meter.val, l_dist_ap_meter.val, l_dist_an_meter.val, l_loss_meter.val, )) else: l_log = '' if cfg.id_loss_weight > 0: id_log = (', idL {:.4f}'.format(id_loss_meter.val)) else: id_log = '' total_loss_log = ', loss {:.4f}'.format(loss_meter.val) log = time_log + \ g_log + l_log + id_log + \ total_loss_log print(log) ############# # Epoch Log # ############# time_log = 'Ep {}, {:.2f}s'.format( ep + 1, time.time() - ep_st, ) if cfg.g_loss_weight > 0: g_log = (', gp {:.2%}, gm {:.2%}, ' 'gd_ap {:.4f}, gd_an {:.4f}, ' 'gL {:.4f}'.format( g_prec_meter.avg, g_m_meter.avg, g_dist_ap_meter.avg, g_dist_an_meter.avg, g_loss_meter.avg, )) else: g_log = '' if cfg.l_loss_weight > 0: l_log = (', lp {:.2%}, lm {:.2%}, ' 'ld_ap {:.4f}, ld_an {:.4f}, ' 'lL {:.4f}'.format( l_prec_meter.avg, l_m_meter.avg, l_dist_ap_meter.avg, l_dist_an_meter.avg, l_loss_meter.avg, )) else: l_log = '' if cfg.id_loss_weight > 0: id_log = (', idL {:.4f}'.format(id_loss_meter.avg)) else: id_log = '' total_loss_log = ', loss {:.4f}'.format(loss_meter.avg) log = time_log + \ g_log + l_log + id_log + \ total_loss_log print(log) # Log to TensorBoard if cfg.log_to_file: if writer is None: writer = SummaryWriter( log_dir=osp.join(cfg.exp_dir, 'tensorboard')) writer.add_scalars( 'loss', dict( global_loss=g_loss_meter.avg, local_loss=l_loss_meter.avg, id_loss=id_loss_meter.avg, loss=loss_meter.avg, ), ep) writer.add_scalars( 'tri_precision', dict( global_precision=g_prec_meter.avg, local_precision=l_prec_meter.avg, ), ep) writer.add_scalars( 'satisfy_margin', dict( global_satisfy_margin=g_m_meter.avg, local_satisfy_margin=l_m_meter.avg, ), ep) writer.add_scalars( 'global_dist', dict( global_dist_ap=g_dist_ap_meter.avg, global_dist_an=g_dist_an_meter.avg, ), ep) writer.add_scalars( 'local_dist', dict( local_dist_ap=l_dist_ap_meter.avg, local_dist_an=l_dist_an_meter.avg, ), ep) # save ckpt if cfg.log_to_file: save_ckpt(modules_optims, ep + 1, 0, cfg.ckpt_file) #if (ep+1)%1==0: # test(load_model_weight=False) # if (ep+1)%10==0: # print(reranking_mAP_list) ######## # Test # ######## test(load_model_weight=False)
def main(): cfg = Config() # Redirect logs to both console and file. if cfg.log_to_file: ReDirectSTD(cfg.log_file, 'stdout', False) ReDirectSTD(cfg.log_err_file, 'stderr', False) TVT, TMO = set_devices(cfg.sys_device_ids) if cfg.seed is not None: set_seed(cfg.seed) # Dump the configurations to log. import pprint pprint.pprint(cfg.__dict__) if cfg.log_to_file: writer = SummaryWriter(log_dir=osp.join(cfg.exp_dir, 'tensorboard')) else: writer = None ########### # Models # ########### model = Model(local_conv_out_channels=cfg.local_conv_out_channels) model_w = get_model_wrapper(model, len(cfg.sys_device_ids) > 1) ############################# # Criteria and Optimizers # ############################# g_tri_loss = TripletLoss(margin=cfg.global_margin) l_tri_loss = TripletLoss(margin=cfg.local_margin) optimizer = optim.Adam(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay) modules_optims = [model, optimizer] ################################ # May Resume Models and Optims # ################################ if cfg.resume: resume_ep, scores = load_ckpt(modules_optims, cfg.ckpt_file) # May Transfer Models and Optims to Specified Device TMO(modules_optims) ########### # Dataset # ########### def feature_func(ims): """A function to be called in the val/test set, to extract features.""" # Set eval mode. # Force all BN layers to use global mean and variance, also disable # dropout. may_set_mode(modules_optims, 'eval') ims = Variable(TVT(torch.from_numpy(ims).float())) global_feat, local_feat = model_w(ims) global_feat = global_feat.data.cpu().numpy() local_feat = local_feat.data.cpu().numpy() return global_feat, local_feat train_set, val_set, test_set = None, None, None if not cfg.only_test: train_set = create_dataset(**cfg.train_set_kwargs) # val_set = create_dataset(**cfg.val_set_kwargs) # val_set.set_feat_func(feature_func) if cfg.only_test or cfg.test: test_set = create_dataset(**cfg.test_set_kwargs) test_set.set_feat_func(feature_func) ######## # Test # ######## if cfg.only_test: print('=====> Test') load_ckpt(modules_optims, cfg.ckpt_file) mAP, cmc_scores, mq_mAP, mq_cmc_scores = test_set.eval( normalize_feat=True, global_weight=cfg.g_test_weight, local_weight=cfg.l_test_weight) return ############ # Training # ############ best_score = scores if cfg.resume else 0 start_ep = resume_ep if cfg.resume else 0 for ep in range(start_ep, cfg.num_epochs): adjust_lr(optimizer, cfg.lr, ep, cfg.num_epochs, cfg.start_decay_epoch) may_set_mode(modules_optims, 'train') epoch_done = False g_prec_meter = AverageMeter() g_m_meter = AverageMeter() g_dist_ap_meter = AverageMeter() g_dist_an_meter = AverageMeter() g_loss_meter = AverageMeter() l_prec_meter = AverageMeter() l_m_meter = AverageMeter() l_dist_ap_meter = AverageMeter() l_dist_an_meter = AverageMeter() l_loss_meter = AverageMeter() loss_meter = AverageMeter() ep_st = time.time() step = 0 while not epoch_done: step += 1 step_st = time.time() ims, im_names, labels, mirrored, epoch_done = train_set.next_batch( ) ims_var = Variable(TVT(torch.from_numpy(ims).float())) labels_t = TVT(torch.from_numpy(labels).long()) global_feat, local_feat = model_w(ims_var) g_loss, p_inds, n_inds, g_dist_ap, g_dist_an = global_loss( g_tri_loss, global_feat, labels_t) if cfg.l_loss_weight == 0: l_loss, l_prec, l_m = 0, 0, 0 else: l_loss, l_dist_ap, l_dist_an = local_loss( l_tri_loss, local_feat, p_inds, n_inds, labels_t) # Let local distance find its own hard samples. # l_loss, l_dist_ap, l_dist_an = local_loss( # l_tri_loss, local_feat, None, None, labels_t) loss = g_loss * cfg.g_loss_weight + l_loss * cfg.l_loss_weight optimizer.zero_grad() loss.backward() optimizer.step() # Step logs # precision g_prec = (g_dist_an > g_dist_ap).data.float().mean() # the proportion of triplets that satisfy margin g_m = (g_dist_an > g_dist_ap + cfg.global_margin).data.float().mean() g_d_ap = g_dist_ap.data.mean() g_d_an = g_dist_an.data.mean() g_prec_meter.update(g_prec) g_m_meter.update(g_m) g_dist_ap_meter.update(g_d_ap) g_dist_an_meter.update(g_d_an) g_loss_meter.update(to_scalar(g_loss)) if cfg.l_loss_weight > 0: # precision l_prec = (l_dist_an > l_dist_ap).data.float().mean() # the proportion of triplets that satisfy margin l_m = (l_dist_an > l_dist_ap + cfg.local_margin).data.float().mean() l_d_ap = l_dist_ap.data.mean() l_d_an = l_dist_an.data.mean() l_prec_meter.update(l_prec) l_m_meter.update(l_m) l_dist_ap_meter.update(l_d_ap) l_dist_an_meter.update(l_d_an) l_loss_meter.update(to_scalar(l_loss)) loss_meter.update(to_scalar(loss)) if step % cfg.log_steps == 0: print( '\tStep {}/Ep {}, {:.2f}s, ' 'gp {:.4f}, gm {:.4f}, gd_ap {:.4f}, gd_an {:.4f}, g_loss {:.4f}, ' 'lp {:.4f}, lm {:.4f}, ld_ap {:.4f}, ld_an {:.4f}, l_loss {:.4f}, ' 'loss: {:.4f}'.format( step, ep + 1, time.time() - step_st, g_prec_meter.val, g_m_meter.val, g_dist_ap_meter.val, g_dist_an_meter.val, g_loss_meter.val, l_prec_meter.val, l_m_meter.val, l_dist_ap_meter.val, l_dist_an_meter.val, l_loss_meter.val, loss_meter.val)) # Epoch logs print( 'Ep {}, {:.2f}s, ' 'gp {:.4f}, gm {:.4f}, gd_ap {:.4f}, gd_an {:.4f}, g_loss {:.4f}, ' 'lp {:.4f}, lm {:.4f}, ld_ap {:.4f}, ld_an {:.4f}, l_loss {:.4f}, ' 'loss: {:.4f}'.format(ep + 1, time.time() - ep_st, g_prec_meter.avg, g_m_meter.avg, g_dist_ap_meter.avg, g_dist_an_meter.avg, g_loss_meter.avg, l_prec_meter.avg, l_m_meter.avg, l_dist_ap_meter.avg, l_dist_an_meter.avg, l_loss_meter.avg, loss_meter.avg)) if cfg.log_to_file: writer.add_scalars( 'loss', dict( global_loss=g_loss_meter.avg, local_loss=l_loss_meter.avg, loss=loss_meter.avg, ), ep) writer.add_scalars( 'tri_precision', dict( global_precision=g_prec_meter.avg, local_precision=l_prec_meter.avg, ), ep) writer.add_scalars( 'satisfy_margin', dict( global_proportion=g_m_meter.avg, local_proportion=l_m_meter.avg, ), ep) writer.add_scalars( 'global_dist', dict( global_dist_ap=g_dist_ap_meter.avg, global_dist_an=g_dist_an_meter.avg, ), ep) writer.add_scalars( 'local_dist', dict( local_dist_ap=l_dist_ap_meter.avg, local_dist_an=l_dist_an_meter.avg, ), ep) mAP = 0 # print('=====> Validation') # mAP, cmc_scores, mq_mAP, mq_cmc_scores = val_set.eval( # normalize_feat=True, # global_weight=cfg.g_test_weight, # local_weight=cfg.l_test_weight) # save ckpt if cfg.save_ckpt: save_ckpt(modules_optims, ep + 1, mAP, cfg.ckpt_file) # if mAP > best_score: # best_score = mAP # shutil.copy(cfg.ckpt_file, cfg.best_ckpt_file) ######## # Test # ######## if cfg.test: print('=====> Test') mAP, cmc_scores, mq_mAP, mq_cmc_scores = test_set.eval( normalize_feat=True, global_weight=cfg.g_test_weight, local_weight=cfg.l_test_weight)