def test(model, test_loader, losses, epoch, nr_query=Config.nr_query): ''' return: cmc1, mAP test model on testset and save result to log. ''' val_start_time = time.time() model.eval() logger.info('testing', 'Start testing') all_features, all_labels, all_cids = [], [], [] history = collections.defaultdict(list) for i, (imgs, labels, cids, types, colors) in tqdm(enumerate(test_loader), desc='testing on epoch-{}'.format(epoch), total=len(test_loader)): imgs, labels, cids = imgs.cuda(), labels.cuda(), cids.cuda() types, colors = types.cuda(), colors.cuda() f_norm, p_type, p_color = model(imgs) triplet_hard_loss = losses['triplet_hard_loss'](f_norm, labels) acc_type = accuracy(p_type, types)[0] acc_color = accuracy(p_color, colors)[0] history['triplet_hard_loss'].append(float(triplet_hard_loss)) history['acc_type'].append(float(acc_type)) history['acc_color'].append(float(acc_color)) all_features.append(f_norm.cpu().detach().numpy()) all_labels.append(labels.cpu().detach().numpy()) all_cids.append(cids.cpu().detach().numpy()) all_features = np.concatenate(all_features, axis=0) all_labels = np.concatenate(all_labels, axis=0) all_cids = np.concatenate(all_cids, axis=0) q_f, g_f = all_features[:nr_query], all_features[nr_query:] q_ids, g_ids = all_labels[:nr_query], all_labels[nr_query:] q_cids, g_cids = all_cids[:nr_query], all_cids[nr_query:] print('Compute CMC and mAP') distance_matrix = get_L2distance_matrix_numpy(q_f, g_f) cmc, mAP = get_cmc_map(distance_matrix, q_ids, g_ids, q_cids, g_cids) val_end_time = time.time() time_spent = sec2min_sec(val_start_time, val_end_time) text = 'Finish testing epoch {:>3}, time spent: [{:>3}mins{:>3}s], performance:\n##'.format( epoch, time_spent[0], time_spent[1]) text += '|CMC1:{:>5.4f} |mAP:{:>5.4f}'.format(cmc[0], mAP) for k, vlist in history.items(): v = float(mean(vlist)) text += '|{}:{:>5.4f} '.format(k, v) logger.add_scalar('TEST/' + k, v, epoch) logger.info('testing', text) logger.add_scalar('TEST/cmc1', cmc[0], epoch) logger.add_scalar('TEST/cmc5', cmc[4], epoch) logger.add_scalar('TEST/cmc10', cmc[9], epoch) logger.add_scalar('TEST/mAP', mAP, epoch) return cmc, mAP
def prepare(args): resume_from_checkpoint = args.resume_from_checkpoint prepare_start_time = time.time() logger.info('global', 'Start preparing.') check_config_dir() logger.info('setting', config_info(), time_report=False) model = Baseline(num_classes=Config.nr_class) logger.info('setting', model_summary(model), time_report=False) logger.info('setting', str(model), time_report=False) train_transforms = transforms.Compose([ transforms.Resize(Config.input_shape), transforms.RandomApply([ transforms.ColorJitter( brightness=0.3, contrast=0.3, saturation=0.3, hue=0) ], p=0.5), transforms.RandomHorizontalFlip(), transforms.Pad(10), transforms.RandomCrop(Config.input_shape), transforms.ToTensor(), transforms.RandomErasing(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_transforms = transforms.Compose([ transforms.Resize(Config.input_shape), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) trainset = Veri776_train(transforms=train_transforms, need_attr=True) testset = Veri776_test(transforms=test_transforms, need_attr=True) pksampler = PKSampler(trainset, p=Config.P, k=Config.K) train_loader = torch.utils.data.DataLoader(trainset, batch_size=Config.batch_size, sampler=pksampler, num_workers=Config.nr_worker, pin_memory=True) test_loader = torch.utils.data.DataLoader( testset, batch_size=Config.batch_size, sampler=torch.utils.data.SequentialSampler(testset), num_workers=Config.nr_worker, pin_memory=True) weight_decay_setting = parm_list_with_Wdecay(model) optimizer = torch.optim.Adam(weight_decay_setting, lr=Config.lr) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_multi_func) losses = {} losses['cross_entropy_loss'] = torch.nn.CrossEntropyLoss() losses['type_ce_loss'] = torch.nn.CrossEntropyLoss() losses['color_ce_loss'] = torch.nn.CrossEntropyLoss() losses['triplet_hard_loss'] = triplet_hard_loss( margin=Config.triplet_margin) for k in losses.keys(): losses[k] = losses[k].cuda() start_epoch = 0 if resume_from_checkpoint and os.path.exists(Config.checkpoint_path): checkpoint = load_checkpoint() start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) scheduler.load_state_dict(checkpoint['scheduler']) # continue training for next the epoch of the checkpoint, or simply start from 1 start_epoch += 1 ret = { 'start_epoch': start_epoch, 'model': model, 'train_loader': train_loader, 'test_loader': test_loader, 'optimizer': optimizer, 'scheduler': scheduler, 'losses': losses } prepare_end_time = time.time() time_spent = sec2min_sec(prepare_start_time, prepare_end_time) logger.info( 'global', 'Finish preparing, time spend: {}mins {}s.'.format( time_spent[0], time_spent[1])) return ret
def start( model, train_loader, test_loader, optimizer, scheduler, losses, start_epoch, ): train_start_time = time.time() best_mAP = 0.0 best_mAP_epoch = 0 best_top1 = 0.0 best_top1_epoch = 0 logger.info('global', 'Start training.') for epoch in range(start_epoch, Config.epoch + 1): train_one_epoch(model, train_loader, losses, optimizer, scheduler, epoch) if epoch % Config.epoch_per_test == 0: cmc, mAP = test(model, test_loader, losses, epoch) top1 = cmc[0] if top1 > best_top1: best_top1 = top1 best_top1_epoch = epoch if mAP > best_mAP: best_mAP = mAP best_mAP_epoch = epoch if epoch % Config.epoch_per_save == 0: if Config.epoch_per_test % Config.epoch_per_save != 0: cmc, mAP = test(model, test_loader, losses, epoch) file_name = 'epoch-{:0>3}'.format(epoch) + '.pth' save_dict = { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), 'top1': cmc[0], 'mAP': mAP } path = os.path.join(Config.model_dir, file_name) logger.info('global', 'Save model to {}'.format(path)) torch.save(save_dict, path) save_dict = { 'epoch': epoch, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict() } save_checkpoint(save_dict) train_end_time = time.time() time_spent = sec2min_sec(train_start_time, train_end_time) text = 'Finish training, time spent: {:>3}mins {:>3}s'.format( time_spent[0], time_spent[1]) logger.info('global', text) text = '##FINISH## best mAP:{:>5.4f} in epoch {:>3}; best top1:{:>5.4f} in epoch{:>3}'.format( best_mAP, best_mAP_epoch, best_top1, best_top1_epoch) logger.info('global', text)
def train_one_epoch(model, train_loader, losses, optimizer, scheduler, epoch): global batch_step epoch_start_time = time.time() logger.info( 'training', 'Start training epoch-{}, lr={:.6}'.format( epoch, get_lr_from_optim(optimizer))) scaler = amp.GradScaler() model.train() history = collections.defaultdict(list) for i, (imgs, labels, _, types, colors) in enumerate(train_loader): batch = i + 1 batch_start_time = time.time() imgs, labels = imgs.cuda(), labels.cuda() types, colors = types.cuda(), colors.cuda() with amp.autocast(): f_bn, p, p_type, p_color = model(imgs) ce_loss = losses['cross_entropy_loss'](p, labels) triplet_hard_loss = losses['triplet_hard_loss'](f_bn, labels) type_ce_loss = losses['type_ce_loss'](p_type, types) color_ce_loss = losses['color_ce_loss'](p_color, colors) loss = Config.weight_ce * ce_loss loss += Config.weight_tri * triplet_hard_loss loss += Config.w_type * type_ce_loss loss += Config.w_color * color_ce_loss scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() optimizer.zero_grad() acc = accuracy(p, labels)[0] acc_type = accuracy(p_type, types)[0] acc_color = accuracy(p_color, colors)[0] batch_end_time = time.time() time_spent = batch_end_time - batch_start_time dist_ap, dist_an = losses['triplet_hard_loss'].get_mean_hard_dist() perform = { 'ce': float(Config.weight_ce * ce_loss), 'tri_h': float(Config.weight_tri * triplet_hard_loss), 'type_ce': float(Config.w_type * type_ce_loss), 'color_ce': float(Config.w_color * color_ce_loss), 'dap_h': float(dist_ap), 'dan_h': float(dist_an), 'acc': float(acc), 'acc_type': float(acc_type), 'acc_color': float(acc_color), 'time(s)': float(time_spent) } if i % Config.batch_per_log == 0: stage = (epoch, batch) text = '' for k, v in perform.items(): text += '|{}:{:<8.4f} '.format(k, float(v)) logger.info('training', text, stage=stage) for k, v in perform.items(): history[k].append(float(v)) if k != 'time(s)': logger.add_scalar('TRAIN_b/' + k, v, batch_step) batch_step += 1 scheduler.step() epoch_end_time = time.time() time_spent = sec2min_sec(epoch_start_time, epoch_end_time) text = 'Finish training epoch {}, time spent: {}mins {}secs, performance:\n##'.format( epoch, time_spent[0], time_spent[1]) for k, vlist in history.items(): v = mean(vlist) text += '|{}:{:>5.4f} '.format(k, v) if k != 'time(s)': logger.add_scalar('TRAIN_e/' + k, v, epoch) logger.info('training', text)
def test(model, test_loader, losses, epoch, nr_query=Config.nr_query): ''' return: cmc1, mAP test model on testset and save result to log. ''' val_start_time = time.time() model.eval() logger.info('testing', 'Start testing') all_features_gpu, all_labels, all_cids, all_mask_gpu = [], [], [], [] all_attn_dist = [] history = collections.defaultdict(list) for i, (imgs, labels, cids) in tqdm(enumerate(test_loader), desc='testing on epoch-{}'.format(epoch), total=len(test_loader)): imgs, labels, cids = imgs.cuda(), labels.cuda(), cids.cuda() f_norm, f_mask = model(imgs) triplet_hard_loss, dist = losses['triplet_hard_loss'](f_norm, f_mask, labels) history['triplet_hard_loss'].append(float(triplet_hard_loss)) all_features_gpu.append(f_norm) all_labels.append(labels.cpu().detach().numpy()) all_cids.append(cids.cpu().detach().numpy()) all_mask_gpu.append(f_mask) all_attn_dist.append(dist.cpu().detach().numpy()) all_features_gpu = torch.cat(all_features_gpu, axis=0) all_features = all_features_gpu.cpu().detach().numpy() all_labels = np.concatenate(all_labels, axis=0) all_cids = np.concatenate(all_cids, axis=0) all_mask_gpu = torch.cat(all_mask_gpu, axis=0) q_f_gpu, g_f_gpu = all_features_gpu[:nr_query], all_features_gpu[nr_query:] q_f, g_f = all_features[:nr_query], all_features[nr_query:] q_ids, g_ids = all_labels[:nr_query], all_labels[nr_query:] q_cids, g_cids = all_cids[:nr_query], all_cids[nr_query:] q_mask, g_mask = all_mask_gpu[:nr_query], all_mask_gpu[nr_query:] print('Compute CMC and mAP') attn_distance_matrix = get_L2distance_matrix_attn_batch( q_f_gpu, g_f_gpu, q_mask, g_mask, temp=Config.temperature) attn_distance_matrix = attn_distance_matrix.cpu().detach().numpy() distance_matrix = get_L2distance_matrix_numpy(q_f, g_f) cmc_a, mAP_a = get_cmc_map(attn_distance_matrix, q_ids, g_ids, q_cids, g_cids) cmc, mAP = get_cmc_map(distance_matrix, q_ids, g_ids, q_cids, g_cids) val_end_time = time.time() time_spent = sec2min_sec(val_start_time, val_end_time) text = 'Finish testing epoch {:>3}, time spent: [{:>3}mins{:>3}s], performance:\n##'.format( epoch, time_spent[0], time_spent[1]) text += 'With ATTENTION> |CMC1:{:>5.4f} |mAP:{:>5.4f} '.format( cmc_a[0], mAP_a) text += 'W/O attention> |CMC1:{:>5.4f} |mAP:{:>5.4f}'.format(cmc[0], mAP) for k, vlist in history.items(): v = float(mean(vlist)) text += '|{}:{:>5.4f} '.format(k, v) logger.add_scalar('TEST/' + k, v, epoch) logger.info('testing', text) logger.add_scalar('TEST/cmc1_a', cmc_a[0], epoch) logger.add_scalar('TEST/cmc5_a', cmc_a[4], epoch) logger.add_scalar('TEST/mAP_a', mAP_a, epoch) logger.add_scalar('TEST/cmc1', cmc[0], epoch) logger.add_scalar('TEST/cmc5', cmc[4], epoch) logger.add_scalar('TEST/mAP', mAP, epoch) return cmc_a, mAP_a, cmc, mAP
def train_one_epoch(model, branches, nr_mask, train_loader, losses, optimizer, scheduler, epoch): global batch_step epoch_start_time = time.time() logger.info( 'training', 'Start training epoch-{}, lr={:.6}'.format( epoch, get_lr_from_optim(optimizer))) scaler = amp.GradScaler() model.train() for branch in branches: branch.train() history = collections.defaultdict(list) for i, (imgs, labels, masks) in enumerate(train_loader): batch = i + 1 batch_start_time = time.time() imgs, labels, masks = imgs.cuda(), labels.cuda(), masks.float().cuda() with amp.autocast(): loss = 0 parsing_celoss = [0] * nr_mask parsing_triloss = [0] * nr_mask x = model(imgs) weights = get_weight(masks) f_main, p = branches[0](x) ce_loss = losses['cross_entropy_loss'][0](p, labels) triplet_hard_loss = losses['triplet_hard_loss'][0](f_main, labels) loss += Config.weight_ce[0] * ce_loss loss += Config.weight_tri[0] * triplet_hard_loss for b, branch in enumerate(branches): if b == 0: # main branch continue mask = masks[:, b - 1:b, ...] f, logit = branch(x, mask) w = weights[:, b - 1] pce_loss = losses['cross_entropy_loss'][b](logit, labels, w) # ptriplet_hard_loss = losses['triplet_hard_loss'][b]( # f, w, labels) # ptriplet_hard_loss = losses['triplet_hard_loss'][0]( # f, labels) parsing_celoss[b - 1] = Config.weight_ce[b] * pce_loss # parsing_triloss += Config.weight_tri[b] * ptriplet_hard_loss loss = loss + sum(parsing_celoss) + sum(parsing_triloss) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() optimizer.zero_grad() acc = accuracy(p, labels)[0] batch_end_time = time.time() time_spent = batch_end_time - batch_start_time dist_ap, dist_an = losses['triplet_hard_loss'][0].get_mean_hard_dist() perform = {} for i in range(nr_mask): if i == 2: continue # skip top-side loss perform.update({'p%d_ce' % (i + 1): float(parsing_celoss[i])}) for i in range(nr_mask): if i == 2: continue # skip top-side loss perform.update({'p%d_tri' % (i + 1): float(parsing_triloss[i])}) perform.update({ 'ce': float(Config.weight_ce[0] * ce_loss), 'tri': float(Config.weight_tri[0] * triplet_hard_loss), 'dap': float(dist_ap), 'dan': float(dist_an), 'acc': float(acc), 't': float(time_spent) }) if i % Config.batch_per_log == 0: stage = (epoch, batch) text = '' for k, v in perform.items(): text += '|{}:{:<6.4f} '.format(k, float(v)) logger.info('training', text, stage=stage) for k, v in perform.items(): history[k].append(float(v)) if k != 'time(s)': logger.add_scalar('TRAIN_b/' + k, v, batch_step) batch_step += 1 scheduler.step() epoch_end_time = time.time() time_spent = sec2min_sec(epoch_start_time, epoch_end_time) text = 'Finish training epoch {}, time spent: {}mins {}secs, performance:\n##'.format( epoch, time_spent[0], time_spent[1]) for k, vlist in history.items(): v = mean(vlist) text += '|{}:{:>5.4f} '.format(k, v) if k != 'time(s)': logger.add_scalar('TRAIN_e/' + k, v, epoch) logger.info('training', text)
def test(model, bclassifier, test_loader, losses, epoch, nr_query=Config.nr_query): ''' return: cmc1, mAP test model on testset and save result to log. ''' val_start_time = time.time() model.eval() bclassifier.eval() logger.info('testing', 'Start testing') all_features_gpu, all_labels, all_cids, all_mask = [], [], [], [] for i, (imgs, labels, cids) in tqdm(enumerate(test_loader), desc='extracting features', total=len(test_loader)): imgs, labels, cids = imgs.cuda(), labels.cuda(), cids.cuda() f_norm, f_mask = model(imgs) all_features_gpu.append(f_norm) all_labels.append(labels) all_mask.append(f_mask) all_cids.append(cids) features = torch.cat(all_features_gpu, axis=0) alllabels = torch.cat(all_labels, axis=0) allmask = torch.cat(all_mask, axis=0) allcids = torch.cat(all_cids, axis=0) q_f_gpu, g_f_gpu = features[:nr_query, ...], features[nr_query:, ...] q_ids, g_ids = alllabels[:nr_query, ...], alllabels[nr_query:, ...] q_cids, g_cids = allcids[:nr_query], allcids[nr_query:] q_mask, g_mask = allmask[:nr_query], allmask[nr_query:] pk_pros = [] for start in tqdm(range(0, q_f_gpu.shape[0], Config.eval_P), desc='computing similarity'): end = min(start + Config.eval_P, q_f_gpu.shape[0]) p_pros = [] for kstart in range(0, g_f_gpu.shape[0], Config.eval_K): kend = min(kstart + Config.eval_K, g_f_gpu.shape[0]) pros = bclassifier(q_f_gpu[start:end, ...], g_f_gpu[kstart:kend, ...]) pros = pros.reshape(end - start, kend - kstart, 2) p_pros.append(pros) p_pros = torch.cat(p_pros, axis=1) pk_pros.append(p_pros) pros = torch.cat(pk_pros, axis=0) pairlabels = idlabel2pairlabel(q_ids, g_ids) acc = accuracy(pros.reshape(-1, 2), pairlabels)[0] distance_matrix = pros[:, :, 0].cpu().detach().numpy() q_ids = q_ids.cpu().detach().numpy() g_ids = g_ids.cpu().detach().numpy() q_cids = q_cids.cpu().detach().numpy() g_cids = g_cids.cpu().detach().numpy() print('Compute CMC and mAP') cmc, mAP = get_cmc_map(distance_matrix, q_ids, g_ids, q_cids, g_cids) val_end_time = time.time() time_spent = sec2min_sec(val_start_time, val_end_time) text = 'Finish testing epoch {:>3}, time spent: [{:>3}mins{:>3}s], performance:\n##'.format( epoch, time_spent[0], time_spent[1]) text += 'W/O attention> |CMC1:{:>5.4f} |mAP:{:>5.4f} |ACC:{:>5.4f} '.format( cmc[0], mAP, acc) logger.info('testing', text) logger.add_scalar('TEST/cmc1', cmc[0], epoch) logger.add_scalar('TEST/cmc5', cmc[4], epoch) logger.add_scalar('TEST/mAP', mAP, epoch) logger.add_scalar('TEST/acc', acc, epoch) return cmc, mAP, acc
def train_one_epoch(model, bclassifier, train_loader, losses, optimizer, scheduler, epoch): global batch_step epoch_start_time = time.time() logger.info( 'training', 'Start training epoch-{}, lr={:.6}'.format( epoch, get_lr_from_optim(optimizer))) scaler = amp.GradScaler() model.train() bclassifier.train() history = collections.defaultdict(list) for i, (imgs, labels) in enumerate(train_loader): batch = i + 1 batch_start_time = time.time() imgs, labels = imgs.cuda(), labels.cuda() pair_labels = idlabel2pairlabel(labels, labels) with amp.autocast(): f_bn, p, f_mask = model(imgs) p_same = bclassifier(f_bn, f_bn) id_loss = losses['cross_entropy_loss'](p, labels) pair_loss = losses['pair_loss'](p_same, pair_labels) loss = Config.weight_ce * id_loss loss += Config.weight_pairloss * pair_loss scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() optimizer.zero_grad() sortmask = torch.sort(f_mask, dim=1)[0] bottom20mean = sortmask[:, :20].mean() top20mean = sortmask[:, -20:].mean() acc = accuracy(p, labels)[0] acc_pair = accuracy(p_same, pair_labels)[0] batch_end_time = time.time() time_spent = batch_end_time - batch_start_time perform = { 'id_loss': float(Config.weight_ce * id_loss), 'pair_loss': float(Config.weight_pairloss * pair_loss), 'acc_id': float(acc), 'acc_pair': float(acc_pair), # 'mtop20': float(top20mean), # 'mbot20': float(bottom20mean), 'time(s)': float(time_spent) } if i % Config.batch_per_log == 0: stage = (epoch, batch) text = '' for k, v in perform.items(): text += '|{}:{:<8.4f} '.format(k, float(v)) logger.info('training', text, stage=stage) for k, v in perform.items(): history[k].append(float(v)) if k != 'time(s)': logger.add_scalar('TRAIN_b/' + k, v, batch_step) batch_step += 1 scheduler.step() epoch_end_time = time.time() time_spent = sec2min_sec(epoch_start_time, epoch_end_time) text = 'Finish training epoch {}, time spent: {}mins {}secs, performance:\n##'.format( epoch, time_spent[0], time_spent[1]) for k, vlist in history.items(): v = mean(vlist) text += '|{}:{:>5.4f} '.format(k, v) if k != 'time(s)': logger.add_scalar('TRAIN_e/' + k, v, epoch) logger.info('training', text)