def optimizer_creation(cfg, model): """ :return: The parameters has been split to 2 groups: ft_params: always ResNet parameters. new_params: always many layers parameters after ResNet. ft_params use cfg.ft_lr new_params use cfg.new_params_lr They are pairs by OrderedDict to save. """ param_groups = model.get_param_groups() optimizer = create_optimizer(param_groups, cfg) recursive_to_device(optimizer.state_dict(), cfg.device) return optimizer
def extract_dataloader_feature(model, dataloader, cfg): """ :param model: :param dataloader: gallery or query dataloader :param cfg: :return: concat_feat_dict_list I concat each batch item together in a dict: The dict should be this format: { 'im_path': [item1, item2, ...], 'feat': [item1, item2, ...], 'label': [item1, item2, ...], 'cam': [item1, item2, ...], } """ feat_dict_list = [] for item in tqdm(dataloader, desc='Extract Feature', miniters=20, ncols=120, unit=' batches'): model.eval() with torch.no_grad(): item = recursive_to_device(item, cfg.device) output = model(item, cfg) part_label_batch = extract_part_label(item, cfg) feat_dict = { 'total_part_label': part_label_batch.cpu().numpy(), 'total_partial_feature': output['pool_feat_list'][1].cpu().numpy(), 'total_pg_global_feature': output['reduction_pool_feat_list'][0].cpu().numpy(), 'total_label': item['label'].cpu().numpy(), 'total_cam': item['cam'].cpu().numpy(), 'list_img': item['im_path'], } feat_dict_list.append(feat_dict) concat_feat_dict = concat_dict_list(feat_dict_list) return concat_feat_dict
def extract_dataloader_feature(model, dataloader, cfg): """ :param model: :param dataloader: gallery or query dataloader :param cfg: :return: concat_feat_dict_list I concat each batch item together in a dict: The dict should be this format: { 'im_path': [item1, item2, ...], 'feat': [item1, item2, ...], 'label': [item1, item2, ...], 'cam': [item1, item2, ...], } """ feat_dict_list = [] for item in tqdm(dataloader, desc='Extract Feature', miniters=20, ncols=120, unit=' batches'): model.eval() with torch.no_grad(): # ff = torch.FloatTensor(item['im'].size(0), 2048).zero_() # for i in range(2): # if i == 1: # item['im'] = fliphor(item['im']) # item = recursive_to_device(item, cfg.device) # output = model(item) # feat = torch.cat(output['reduction_pool_feat_list'], 1) # # feat = feat.cpu().numpy() # feat = feat.cpu() # ff = ff + feat item = recursive_to_device(item, cfg.device) output = model(item, cfg) _, D = output['reduction_pool_feat_list'][0].size() feat_list_length = len(output['reduction_pool_feat_list']) ff_length = feat_list_length * D ff = torch.FloatTensor(item['im'].size(0), ff_length).zero_() # output['normalize_reduction_pool_feat_list'] = [normalize(feat) for feat in output['reduction_pool_feat_list']] feat = torch.cat( output['reduction_pool_feat_list'][:feat_list_length], 1) # feat = feat.cpu().numpy() feat = feat.cpu() fnorm = torch.norm(feat, p=2, dim=1, keepdim=True) ff = feat.div(fnorm.expand_as(ff)).numpy() feat_dict = { 'im_path': item['im_path'], # 'feat': feat, 'feat': ff, 'list_img': item['im_path'] } if 'label' in item: feat_dict['label'] = item['label'].cpu().numpy() if 'cam' in item: feat_dict['cam'] = item['cam'].cpu().numpy() feat_dict_list.append(feat_dict) concat_feat_dict = concat_dict_list(feat_dict_list) return concat_feat_dict
def extract_dataloader_feature(model, dataloader, cfg, use_gcn): """ :param model: :param dataloader: gallery or query dataloader :param cfg: :param use_gcn: feat is if gcned or not :return: concat_feat_dict_list I concat each batch item together in a dict: The dict should be this format: { 'im_path': [item1, item2, ...], 'feat': [item1, item2, ...], 'label': [item1, item2, ...], 'cam': [item1, item2, ...], } """ feat_dict_list = [] for item in tqdm(dataloader, desc='Extract Feature', miniters=20, ncols=120, unit=' batches'): model.eval() with torch.no_grad(): item = recursive_to_device(item, cfg.device) output = model(item, cfg) N, keypoints_num = output['keypoints_confidence'].shape output['keypoints_confidence'] = torch.sqrt( output['keypoints_confidence']).unsqueeze(2).repeat( [1, 1, 2048]).view([N, 2048 * keypoints_num]) feat_stage1 = output['keypoints_confidence'] * torch.cat( output['bned_feat_vec_list'], dim=1) feat_stage2 = torch.cat( [i.unsqueeze(1) for i in output['bned_feat_vec_list']], dim=1) gcned_feat_stage1 = output['keypoints_confidence'] * torch.cat( output['bned_gcned_feat_vec_list'], dim=1) gcned_feat_stage2 = torch.cat( [i.unsqueeze(1) for i in output['bned_gcned_feat_vec_list']], dim=1) if use_gcn: feat_stage1 = gcned_feat_stage1 feat_stage2 = gcned_feat_stage2 else: feat_stage1 = feat_stage1 feat_stage2 = feat_stage2 feat_dict = { 'feat_stage1': feat_stage1.cpu().numpy(), 'feat_stage2': feat_stage2.cpu().numpy(), 'label': item['label'].cpu().numpy(), 'cam': item['cam'].cpu().numpy(), 'list_img': item['im_path'], } feat_dict_list.append(feat_dict) concat_feat_dict = concat_dict_list(feat_dict_list) return concat_feat_dict
def train(self): # dataset_sizes = len(self.source_train_loader.dataset) # class_num = self.source_train_loader.dataset.num_ids # assert self.cfg.model.num_classes == class_num, "cfg.model.num_classes should be {} in create_train_dataloader.py".format(class_num) print("End epoch is:", self.cfg.optim.epochs) for epoch in range(self.current_ep, self.cfg.optim.epochs): self.epoch_start_time = time.time() self.model.set_train_mode( fix_ft_layers=self.cfg.optim.phase == 'pretrain') for index, item in enumerate(self.source_train_loader): if self.lr_scheduler is not None: self.lr_scheduler.step() self.optimizer.zero_grad() # Source item item = recursive_to_device(item, self.cfg.device) pred = self.model.forward(item, cfg=cfg, forward_type='Supervised') loss = 0 for loss_type in [ self.cfg.id_loss, self.cfg.tri_loss, self.cfg.pgfa_loss, self.cfg.src_ps_loss, self.cfg.src_psgp_loss ]: if loss_type.use is True: loss += self.loss_functions[loss_type.name]( item, pred, step=self.current_step)['loss'] if isinstance(loss, torch.Tensor): loss.backward() self.optimizer.step() if ((self.current_step + 1) % self.steps_per_log == 0) and (self.print_step_log is not None): self.print_step_log(self.cfg, self.current_ep, self.current_step, self.optimizer, self.loss_functions, self.analyze_functions, self.epoch_start_time) self.current_step += 1 if (self.trial_run_steps is not None) and (index + 1 >= self.trial_run_steps): break if ((self.current_ep + 1) % self.eps_per_log == 0) and (self.print_ep_log is not None): self.print_ep_log() self.current_ep += 1 score_str = self.may_test() self.may_save_ckpt(score_str, self.current_ep)
def visualization(model, dataloader, cfg): if cfg.vis.heat_map.use is True: savename = cfg.vis.heat_map.save_dir + '/' + dataloader.dataset.authority if not osp.exists(savename): os.makedirs(savename) print(savename, "has been created.") for item in tqdm(dataloader, desc='Extract Feature', miniters=20, ncols=120, unit=' batches'): model.eval() with torch.no_grad(): item = recursive_to_device(item, cfg.device) output = model(item) draw_heat_map(output['ps_pred_list'], im_path=item['im_path'], savename=savename) return
def train(self): print("End epoch is:", self.cfg.optim.epochs) for epoch in range(self.current_ep, self.cfg.optim.epochs): self.epoch_start_time = time.time() self.model.set_train_mode( fix_ft_layers=self.cfg.optim.phase == 'pretrain') # Unsupervised loader iter if self.cfg.dataset.train.type is 'Unsupervised': target_iter = iter(self.target_train_loader) # Train for index, item in enumerate(self.source_train_loader): if self.lr_scheduler is not None: self.lr_scheduler.step() self.optimizer.zero_grad() # Source item item = recursive_to_device(item, self.cfg.device) # Target item if self.cfg.dataset.train.type is 'Unsupervised': try: item_target = next(target_iter) except: target_iter = iter(self.target_train_loader) item_target = next(target_iter) item_target = recursive_to_device(item_target, self.cfg.device) # Supervised Part pred = self.model.forward(item, cfg=self.cfg, forward_type='Supervised') source_loss = 0 for loss_type in [ self.cfg.id_loss, self.cfg.tri_loss, self.cfg.src_multi_seg_loss ]: if loss_type.use is True: source_loss += self.loss_functions[loss_type.name]( item, pred, step=self.current_step)['loss'] if self.cfg.dataset.train.type is 'Unsupervised': # Unsupervised Part pred = self.model.forward(item_target, forward_type='Unsupervised') un_loss = self.loss_functions[self.cfg.inv_loss.name]( pred, epoch=epoch, step=self.current_step)['loss'] loss = (1 - self.cfg.inv_loss.lmd ) * source_loss + self.cfg.inv_loss.lmd * un_loss else: loss = source_loss if isinstance(loss, torch.Tensor): loss.backward() self.optimizer.step() if ((self.current_step + 1) % self.steps_per_log == 0) and (self.print_step_log is not None): self.print_step_log(self.cfg, self.current_ep, self.current_step, self.optimizer, self.loss_functions, self.analyze_functions, self.epoch_start_time) self.current_step += 1 if (self.trial_run_steps is not None) and (index + 1 >= self.trial_run_steps): break if ((self.current_ep + 1) % self.eps_per_log == 0) and (self.print_ep_log is not None): self.print_ep_log() self.current_ep += 1 score_str = self.may_test() self.may_save_ckpt(score_str, self.current_ep)