def make_test_loader_M_C_D(root, name): query_dataset = Market1501(root + '/query', transform=test_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) query_flip_dataset = Market1501(root + '/query', transform=test_flip_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) query_loader = DataLoader(query_dataset, batch_size=batch_test, num_workers=num_workers) query_flip_loader = DataLoader(query_flip_dataset, batch_size=batch_test, num_workers=num_workers) test_dataset = Market1501(root + '/bounding_box_test', transform=test_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) test_flip_dataset = Market1501(root + '/bounding_box_test', transform=test_flip_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) test_loader = DataLoader(test_dataset, batch_size=batch_test, num_workers=num_workers) test_flip_loader = DataLoader(test_flip_dataset, batch_size=batch_test, num_workers=num_workers) return query_loader, query_flip_loader, test_loader, test_flip_loader
def __init__(self, split, dataloader): print("[*] Loading Market1501") market = Market1501('gt_bbox', **dataloader) print("[*] Loading CUHK03") cuhk = CUHK03('labeled', **dataloader) print("[*] Loading MOT") mot = MOTreIDWrapper(split, dataloader) self.dataset = ConcatDataset([market, cuhk, mot])
def read_and_prepare_data(): dataset = Market1501(FLAGS.inputs_path, num_validation_y=0.1, seed=1234) X_train, Y_train, _ = dataset.read_train() X_train = np.array(X_train) Y_train = np.array(Y_train) X_val, Y_val, _ = dataset.read_validation() X_val = np.array(X_val) Y_val = np.array(Y_val) return X_train, Y_train, X_val, Y_val
def train(): ## data logger.info('creating dataloader') dataset = Market1501('./dataset/Market-1501-v15.09.15/bounding_box_train', is_train=True) num_classes = dataset.get_num_classes() sampler = BalancedSampler(dataset, 16, 4) dl = DataLoader(dataset, batch_sampler=sampler, num_workers=8) ## network and loss logger.info('setup model and loss') sphereloss = SphereLoss(1024, num_classes) sphereloss.cuda() net = Network_D() net = nn.DataParallel(net) net.train() net.cuda() ## optimizer logger.info('creating optimizer') params = list(net.parameters()) params += list(sphereloss.parameters()) optim = torch.optim.Adam(params, lr=1e-3) ## training logger.info('start training') t_start = time.time() loss_it = [] for ep in range(150): optim, lrs = lr_scheduler(ep, optim) for it, (imgs, lbs, ids) in enumerate(dl): imgs = imgs.cuda() lbs = lbs.cuda() embs = net(imgs) loss = sphereloss(embs, lbs) optim.zero_grad() loss.backward() optim.step() loss_it.append(loss.detach().cpu().numpy()) if it % 10 == 0 and it != 0: t_end = time.time() t_interval = t_end - t_start log_loss = sum(loss_it) / len(loss_it) msg = 'epoch: {}, iter: {}, loss: {:4f}, lr: {}, time: {:4f}'.format( ep, it, log_loss, lrs, t_interval) logger.info(msg) loss_it = [] t_start = t_end ## save model torch.save(net.module.state_dict(), './res/model_final.pkl') logger.info('\nTraining done, model saved to {}\n\n'.format( './res/model_final.pkl'))
def setUp(self): self.batch_id = 4 self.batch_image = 16 self.data_source = Market1501(root + '/bounding_box_train', transform=ToTensor()) self.sampler = RandomIdSampler(self.data_source, batch_image=self.batch_image) self.data_loader = DataLoader(self.data_source, sampler=self.sampler, batch_size=self.batch_id * self.batch_image)
def test_test(self): test = Market1501(root + '/bounding_box_test') self.assertEqual(15913, len(test.imgs)) self.assertEqual(15913, len(test.ids)) self.assertEqual(15913, len(test.cameras)) self.assertEqual(751, len(test._id2label)) self.assertTrue(test.imgs[0].endswith('0000_c1s1_000151_01.jpg')) self.assertEqual(0, test.ids[0]) self.assertEqual(1, test.cameras[0]) self.assertEqual(0, test._id2label[0]) img, target = test[0] self.assertEqual(0, target) self.assertTrue(test.imgs[2798].endswith('0001_c1s1_001051_03.jpg')) self.assertEqual(1, test.ids[2798]) self.assertEqual(1, test.cameras[2798]) self.assertEqual(1, test._id2label[1]) img, target = test[2798] self.assertEqual(1, target)
def test_query(self): query = Market1501(root + '/query') self.assertEqual(3368, len(query.imgs)) self.assertEqual(3368, len(query.ids)) self.assertEqual(3368, len(query.cameras)) self.assertEqual(750, len(query._id2label)) self.assertTrue(query.imgs[0].endswith('0001_c1s1_001051_00.jpg')) self.assertEqual(1, query.ids[0]) self.assertEqual(1, query.cameras[0]) self.assertEqual(0, query._id2label[1]) _, target = query[0] self.assertEqual(0, target) self.assertTrue(query.imgs[6].endswith('0003_c1s6_015971_00.jpg')) self.assertEqual(3, query.ids[6]) self.assertEqual(1, query.cameras[6]) self.assertEqual(1, query._id2label[3]) _, target = query[6] self.assertEqual(1, target)
def test_train(self): train = Market1501(root + '/bounding_box_train') self.assertEqual(12936, len(train.imgs)) self.assertEqual(12936, len(train.ids)) self.assertEqual(12936, len(train.cameras)) self.assertEqual(751, len(train._id2label)) self.assertTrue(train.imgs[0].endswith('0002_c1s1_000451_03.jpg')) self.assertEqual(2, train.ids[0]) self.assertEqual(1, train.cameras[0]) self.assertEqual(0, train._id2label[2]) _, target = train[0] self.assertEqual(0, target) self.assertTrue(train.imgs[46].endswith('0007_c1s6_028546_01.jpg')) self.assertEqual(7, train.ids[46]) self.assertEqual(1, train.cameras[46]) self.assertEqual(1, train._id2label[7]) _, target = train[46] self.assertEqual(1, target)
def embed(): ## load checkpoint res_pth = './res' mod_pth = osp.join(res_pth, 'model_final.pkl') net = Network_D() net.load_state_dict(torch.load(mod_pth)) net.cuda() net.eval() ## data loader query_set = Market1501('./dataset/Market-1501-v15.09.15/query', is_train=False) gallery_set = Market1501( './dataset/Market-1501-v15.09.15/bounding_box_test', is_train=False) query_loader = DataLoader(query_set, batch_size=32, num_workers=4, drop_last=False) gallery_loader = DataLoader(gallery_set, batch_size=32, num_workers=4, drop_last=False) ## embed logger.info('embedding query set ...') query_pids = [] query_camids = [] query_embds = [] for i, (im, _, ids) in enumerate(tqdm(query_loader)): embds = [] for crop in im: crop = crop.cuda() embds.append(net(crop).detach().cpu().numpy()) embed = sum(embds) / len(embds) pid = ids[0].numpy() camid = ids[1].numpy() query_embds.append(embed) query_pids.extend(pid) query_camids.extend(camid) query_embds = np.vstack(query_embds) query_pids = np.array(query_pids) query_camids = np.array(query_camids) logger.info('embedding gallery set ...') gallery_pids = [] gallery_camids = [] gallery_embds = [] for i, (im, _, ids) in enumerate(tqdm(gallery_loader)): embds = [] for crop in im: crop = crop.cuda() embds.append(net(crop).detach().cpu().numpy()) embed = sum(embds) / len(embds) pid = ids[0].numpy() camid = ids[1].numpy() gallery_embds.append(embed) gallery_pids.extend(pid) gallery_camids.extend(camid) gallery_embds = np.vstack(gallery_embds) gallery_pids = np.array(gallery_pids) gallery_camids = np.array(gallery_camids) ## dump embeds results embd_res = (query_embds, query_pids, query_camids, gallery_embds, gallery_pids, gallery_camids) with open('./res/embds.pkl', 'wb') as fw: pickle.dump(embd_res, fw) logger.info('embedding done, dump to: ./res/embds.pkl') return embd_res
def run(): batch_id = 16 batch_image = 4 batch_test = 32 train_transform = transforms.Compose([ transforms.Resize((384, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) train_dataset = Market1501(args.root + '/bounding_box_train', transform=train_transform) train_loader = dataloader.DataLoader(train_dataset, sampler=RandomIdSampler(train_dataset, batch_image=batch_image), batch_size=batch_id * batch_image, num_workers=args.workers) test_transform = transforms.Compose([ transforms.Resize((384, 128)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_flip_transform = transforms.Compose([ transforms.Resize((384, 128)), functional.hflip, transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) query_dataset = Market1501(args.root + '/query', transform=test_transform) query_flip_dataset = Market1501(args.root + '/query', transform=test_flip_transform) query_loader = dataloader.DataLoader(query_dataset, batch_size=batch_test, num_workers=args.workers) query_flip_loader = dataloader.DataLoader(query_flip_dataset, batch_size=batch_test, num_workers=args.workers) test_dataset = Market1501(args.root + '/bounding_box_test', transform=test_transform) test_flip_dataset = Market1501(args.root + '/bounding_box_test', transform=test_flip_transform) test_loader = dataloader.DataLoader(test_dataset, batch_size=batch_test, num_workers=args.workers) test_flip_loader = dataloader.DataLoader(test_flip_dataset, batch_size=batch_test, num_workers=args.workers) mgn = MGN(num_classes=len(train_dataset.unique_ids)).to(DEVICE) cross_entropy_loss = nn.CrossEntropyLoss() triplet_semihard_loss = TripletSemihardLoss(margin=1.2) optimizer = optim.SGD(mgn.parameters(), lr=1e-2, momentum=0.9, weight_decay=5e-4) scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [40, 60], gamma=0.1) epochs = 80 for epoch in range(epochs): mgn.train() scheduler.step() running_loss = 0.0 for i, data in enumerate(train_loader): inputs, labels = data inputs, labels = inputs.to(DEVICE), labels.to(DEVICE) optimizer.zero_grad() outputs = mgn(inputs) losses = [triplet_semihard_loss(output, labels) for output in outputs[1]] + \ [cross_entropy_loss(output, labels) for output in outputs[2]] loss = sum(losses) / len(losses) loss.backward() optimizer.step() running_loss += loss.item() print('%d/%d - %d/%d - loss: %f' % (epoch + 1, epochs, i, len(train_loader), loss.item())) print('epoch: %d/%d - loss: %f' % (epoch + 1, epochs, running_loss / len(train_loader))) if epoch % 10 == 9: mgn.eval() query = np.concatenate([mgn(inputs.to(DEVICE))[0].detach().cpu().numpy() for inputs, _ in query_loader]) query_flip = np.concatenate([mgn(inputs.to(DEVICE))[0].detach().cpu().numpy() for inputs, _ in query_flip_loader]) test = np.concatenate([mgn(inputs.to(DEVICE))[0].detach().cpu().numpy() for inputs, _ in test_loader]) test_flip = np.concatenate([mgn(inputs.to(DEVICE))[0].detach().cpu().numpy() for inputs, _ in test_flip_loader]) dist = cdist((query + query_flip) / 2., (test + test_flip) / 2.) # dist = cdist(normalize(query + query_flip), normalize(test + test_flip)) r = cmc(dist, query_dataset.ids, test_dataset.ids, query_dataset.cameras, test_dataset.cameras, separate_camera_set=False, single_gallery_shot=False, first_match_break=True) m_ap = mean_ap(dist, query_dataset.ids, test_dataset.ids, query_dataset.cameras, test_dataset.cameras) print('epoch[%d]: mAP=%f, r@1=%f, r@3=%f, r@5=%f, r@10=%f' % (epoch + 1, m_ap, r[0], r[2], r[4], r[9]))
def train(): ## data P, K = 16, 8 batchsize = P * K logger.info('creating dataloader') dataset = Market1501('./dataset/Market-1501-v15.09.15/bounding_box_train', is_train=True) num_classes = dataset.get_num_classes() sampler = BalancedSampler(dataset, P, K) dl = DataLoader(dataset, batch_sampler=sampler, num_workers=8) ## network and loss logger.info('setup model and loss') bottleneck_loss = BottleneckLoss(2048, num_classes) bottleneck_loss.cuda() bottleneck_loss.train() net = Embeddor() net.cuda() net.train() # net = nn.DataParallel(net) ## optimizer logger.info('creating optimizer') lr = 0.1 momentum = 0.9 params = list(net.parameters()) params += list(bottleneck_loss.parameters()) optim = torch.optim.SGD(params, lr=lr, momentum=momentum) ## training logger.info('start training') n_epochs = 140 t_start = time.time() loss_it = [] for ep in range(n_epochs): optim, lrs = lr_scheduler(ep, optim) for it, (imgs, lbs, _) in enumerate(dl): imgs = imgs.cuda() lbs = lbs.cuda() optim.zero_grad() embs_org, embs_sft = net(imgs) loss_org = bottleneck_loss(embs_org, lbs) loss_sft = bottleneck_loss(embs_sft, lbs) loss = loss_org + loss_sft loss.backward() optim.step() loss = loss.cpu().item() loss_it.append(loss) # print logging message t_end = time.time() t_interval = t_end - t_start log_loss = sum(loss_it) / len(loss_it) msg = ', '.join( ['epoch: {}', 'loss: {:.4f}', 'lr: {}', 'time: {:.4f}']).format(ep, log_loss, lrs, t_interval) logger.info(msg) loss_it = [] t_start = t_end ## save model if hasattr(net, 'module'): state_dict = net.module.state_dict() else: state_dict = net.state_dict() torch.save(state_dict, './res/model_final.pth') logger.info('\nTraining done, model saved to {}\n\n'.format( './res/model_final.pth'))
from market1501 import Market1501 from transform import Train_Transform, Val_Transform from torch.utils.data import DataLoader from model import Train_Model from tripletloss import TripletLoss from cross_entropy_smooth import CrossEntropySmooth from optimizer import Make_Optimizer, Warmup import os import time if torch.cuda.is_available(): print("Support GPU!") device = torch.device('cuda') else: device = torch.device('cpu') market = Market1501(root='./') train_transform = Train_Transform(True) val_transform = Val_Transform() train_sampler = RandomIdentitySampler(market.train, 16, 4) train_dataset = ImageDataset(dataset=market.train, transform=train_transform) val_dataset = ImageDataset(dataset=market.test + market.query, transform=val_transform) train_dataloader = DataLoader(train_dataset, 64, False, train_sampler) val_dataloader = DataLoader(val_dataset, 128, False) Model = Train_Model().to(device) IDloss = CrossEntropySmooth(market.num_train_id) optimizer = Make_Optimizer(Model, 3.5e-5) tripletloss = TripletLoss(0.3) warmup = Warmup(optimizer) EPOCH = 120 for epoch in range(EPOCH):
def run(args): gpuId, epochs, weight_decay, batch_id, batch_image, lr_1, lr_2, erasing_p, sampling, exp_dir, trainset_name, cd_trainset_name, testset_names, rand_crop, head_1part_stride = \ args.gpuId, args.epochs, args.weight_decay, args.batch_id, args.batch_image, args.lr_1, args.lr_2, args.erasing_p, args.sampling, args.exp_dir, args.trainset_name, args.cd_trainset_name, args.testset_names, args.rand_crop, args.head_1part_stride DEVICE = torch.device("cuda:" + gpuId if torch.cuda.is_available() else "cpu") print(DEVICE) num_workers = 4 batch_test = 64 #32 train_list = [ transforms.Resize((400, 144)), transforms.RandomCrop((384, 128)) ] if rand_crop else [transforms.Resize((384, 128))] train_list += [ transforms.ToTensor(), ] if erasing_p > 0: train_list = train_list + [ RandomErasing(probability=erasing_p, mean=[0.0, 0.0, 0.0]) ] train_list += [ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ] train_transform = transforms.Compose(train_list) if trainset_name in ['market1501', 'cuhk03', 'duke']: root = get_dataset_root(trainset_name) train_dataset = Market1501( root + '/bounding_box_train', transform=train_transform, training=True, kpt_file=trainset_name + '-kpt.pkl' if args.pap else None, ps_dir=root + '_ps_label' if args.src_ps_lw > 0 else None) elif trainset_name in ['msmt17']: train_dataset = MSMT17(transform=train_transform, training=True, use_kpt=args.pap, use_ps=args.src_ps_lw > 0, split='train') else: raise ValueError('Invalid train set {}'.format(trainset_name)) train_loader = DataLoader(train_dataset, sampler=RandomIdSampler(train_dataset, batch_image=batch_image), batch_size=batch_id * batch_image, num_workers=num_workers, drop_last=True) if args.cd_ps_lw > 0: if cd_trainset_name in ['market1501', 'cuhk03', 'duke']: cd_train_dataset = Market1501(get_dataset_root(cd_trainset_name) + '/bounding_box_train', transform=train_transform, training=True, kpt_file=None, ps_dir=cd_trainset_name + '-ps') elif cd_trainset_name in ['msmt17']: cd_train_dataset = MSMT17(transform=train_transform, training=True, use_kpt=False, use_ps=True) else: raise ValueError( 'Invalid cd train set {}'.format(cd_trainset_name)) cd_train_loader = InfiniteNextBatch( DataLoader(cd_train_dataset, batch_size=args.cd_train_batch_size, num_workers=num_workers, drop_last=True)) test_transform = transforms.Compose([ transforms.Resize((384, 128)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_flip_transform = transforms.Compose([ transforms.Resize((384, 128)), functional.hflip, transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) def make_test_loader_M_C_D(root, name): query_dataset = Market1501(root + '/query', transform=test_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) query_flip_dataset = Market1501(root + '/query', transform=test_flip_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) query_loader = DataLoader(query_dataset, batch_size=batch_test, num_workers=num_workers) query_flip_loader = DataLoader(query_flip_dataset, batch_size=batch_test, num_workers=num_workers) test_dataset = Market1501(root + '/bounding_box_test', transform=test_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) test_flip_dataset = Market1501(root + '/bounding_box_test', transform=test_flip_transform, training=False, kpt_file=name + '-kpt.pkl' if args.pap else None) test_loader = DataLoader(test_dataset, batch_size=batch_test, num_workers=num_workers) test_flip_loader = DataLoader(test_flip_dataset, batch_size=batch_test, num_workers=num_workers) return query_loader, query_flip_loader, test_loader, test_flip_loader def make_test_loader_MS_PR_PI(name): if name == 'msmt17': dclass = MSMT17 elif name == 'partial_reid': dclass = PartialREID elif name == 'partial_ilids': dclass = PartialiLIDs else: raise ValueError('Invalid dataset name {}'.format(name)) q_set = dclass(transform=test_transform, training=False, use_kpt=args.pap, use_ps=False, split='query') q_flip_set = dclass(transform=test_flip_transform, training=False, use_kpt=args.pap, use_ps=False, split='query') q_loader = DataLoader(q_set, batch_size=batch_test, num_workers=num_workers) q_flip_loader = DataLoader(q_flip_set, batch_size=batch_test, num_workers=num_workers) g_set = dclass(transform=test_transform, training=False, use_kpt=args.pap, use_ps=False, split='gallery') g_flip_set = dclass(transform=test_flip_transform, training=False, use_kpt=args.pap, use_ps=False, split='gallery') g_loader = DataLoader(g_set, batch_size=batch_test, num_workers=num_workers) g_flip_loader = DataLoader(g_flip_set, batch_size=batch_test, num_workers=num_workers) return q_loader, q_flip_loader, g_loader, g_flip_loader def make_test_loader(name): if name in ['market1501', 'cuhk03', 'duke']: return make_test_loader_M_C_D(get_dataset_root(name), name) elif name in ['msmt17', 'partial_reid', 'partial_ilids']: return make_test_loader_MS_PR_PI(name) test_loaders = [make_test_loader(name) for name in testset_names] mgn = MGN(len(train_dataset.unique_ids), args) if torch.cuda.device_count() > 1: mgn = nn.DataParallel(mgn) mgn = mgn.to(DEVICE) vanilla_cross_entropy_loss = nn.CrossEntropyLoss() cross_entropy_loss = nn.CrossEntropyLoss(reduce=False) triplet_semihard_loss = TripletSemihardLoss( margin=0.1, DEVICE=DEVICE, sampling=sampling, batch_id=batch_id, batch_image=batch_image) #batch_hard, .'curriculum' ps_loss = PSLoss() optimizer_start1 = optim.SGD(mgn.parameters(), lr=lr_1, momentum=0.9, weight_decay=weight_decay) optimizer_start2 = optim.SGD(mgn.parameters(), lr=lr_2, momentum=0.9, weight_decay=weight_decay) scheduler_1 = optim.lr_scheduler.MultiStepLR(optimizer_start1, [140, 180], gamma=0.1) scheduler_2 = optim.lr_scheduler.MultiStepLR( optimizer_start2, [140, 180], gamma=0.1) # best [140, 180] [120, 160] def get_model_input(inputs, target): dic = {'im': inputs.to(DEVICE)} if 'pap_mask_2p' in target: dic['pap_mask_2p'] = target['pap_mask_2p'].to(DEVICE) dic['pap_mask_3p'] = target['pap_mask_3p'].to(DEVICE) return dic def extract_loader_feat(loader, verbose=False): feat = [] vis = [] i = 0 for inputs, target in loader: if verbose: print(i) i += 1 with torch.no_grad(): output = mgn(get_model_input(inputs, target)) feat.append(output[1].detach().cpu().numpy()) if args.pap: vis_ = np.concatenate([ np.ones([len(output[1]), 3]), torch.stack(output[5 + 3 + 5:5 + 3 + 5 + 5], 1).detach().cpu().numpy() ], 1) vis.append(vis_) feat = np.concatenate(feat) vis = np.concatenate(vis) if args.pap else None return feat, vis def test(query_loader, query_flip_loader, test_loader, test_flip_loader, trainset_name, testset_name, epoch, verbose=False): cache_file = '{}/feat_cache-{}_to_{}.pkl'.format( exp_dir, trainset_name, testset_name) if args.use_feat_cache: assert os.path.exists( cache_file), "Feature cache file {} does not exist!".format( cache_file) query_2, q_vis, query_flip_2, q_vis, test_2, test_vis, test_flip_2, test_vis, q_ids, q_cams, g_ids, g_cams = load_pickle( cache_file) else: query_2, q_vis = extract_loader_feat(query_loader, verbose=verbose) query_flip_2, q_vis = extract_loader_feat(query_flip_loader, verbose=verbose) test_2, test_vis = extract_loader_feat(test_loader, verbose=verbose) test_flip_2, test_vis = extract_loader_feat(test_flip_loader, verbose=verbose) q_ids = query_loader.dataset.ids q_cams = query_loader.dataset.cameras g_ids = test_loader.dataset.ids g_cams = test_loader.dataset.cameras save_pickle([ query_2, q_vis, query_flip_2, q_vis, test_2, test_vis, test_flip_2, test_vis, q_ids, q_cams, g_ids, g_cams ], cache_file) if args.test_which_feat > 0: # TODO: implement for pap idx = args.test_which_feat query_2 = query_2[:, 256 * idx - 256:256 * idx] query_flip_2 = query_flip_2[:, 256 * idx - 256:256 * idx] test_2 = test_2[:, 256 * idx - 256:256 * idx] test_flip_2 = test_flip_2[:, 256 * idx - 256:256 * idx] query = normalize(query_2 + query_flip_2) test = normalize(test_2 + test_flip_2) if verbose: print('query.shape:', query.shape) print('test.shape:', test.shape) if args.pap: print('q_vis.shape:', q_vis.shape) print('test_vis.shape:', test_vis.shape) if args.pap: dist_1 = compute_dist_with_visibility(query, test, q_vis, test_vis, dist_type='euclidean', avg_by_vis_num=False) else: dist_1 = cdist(query, test) r_1 = cmc(dist_1, q_ids, g_ids, q_cams, g_cams, separate_camera_set=False, single_gallery_shot=False, first_match_break=True) m_ap_1 = mean_ap(dist_1, q_ids, g_ids, q_cams, g_cams) print('EPOCH [%d] %s -> %s: mAP=%f, r@1=%f, r@3=%f, r@5=%f, r@10=%f' % (epoch + 1, trainset_name, testset_name, m_ap_1, r_1[0], r_1[2], r_1[4], r_1[9])) if args.only_test: mgn.eval() if not args.use_feat_cache: if args.model_weight_file: model_weight_file = args.model_weight_file else: model_weight_file = '{}/model_weight.pth'.format(exp_dir) load_model_weight((mgn.module if hasattr(mgn, 'module') else mgn), model_weight_file) for name, test_loader in zip(testset_names, test_loaders): test(test_loader[0], test_loader[1], test_loader[2], test_loader[3], trainset_name, name, -1, verbose=False) exit() for epoch in range(epochs): mgn.train() scheduler_1.step() scheduler_2.step() running_loss = 0.0 running_loss_1 = 0.0 running_loss_2 = 0.0 if epoch < 20: optimizer_1 = optim.SGD(mgn.parameters(), lr=0.01 + 0.0045 * epoch, momentum=0.9, weight_decay=weight_decay) optimizer_2 = optim.SGD(mgn.parameters(), lr=0.001 + 0.00045 * epoch, momentum=0.9, weight_decay=weight_decay) else: optimizer_1 = optimizer_start1 optimizer_2 = optimizer_start2 for i, data in enumerate(train_loader): inputs, target = data inputs = inputs.to(DEVICE) for k, v in target.items(): target[k] = v.to(DEVICE) labels = target['id'] outputs = mgn(get_model_input(inputs, target)) optimizer_1.zero_grad() if args.pap: losses_1 = [ vanilla_cross_entropy_loss(output, labels) for output in outputs[5:5 + 3] ] + [(cross_entropy_loss(output, labels) * v).sum() / (v.sum() + 1e-12) for output, v in zip(outputs[5 + 3:5 + 3 + 5], outputs[5 + 3 + 5:5 + 3 + 5 + 5])] else: losses_1 = [ vanilla_cross_entropy_loss(output, labels) for output in outputs[5:5 + 8] ] loss_1 = sum(losses_1) / len(losses_1) psl = 0 if args.src_ps_lw > 0: psl = (ps_loss(outputs[-3], target['ps_label']) + ps_loss(outputs[-2], target['ps_label']) + ps_loss(outputs[-1], target['ps_label'])) / 3. (loss_1 + psl * args.src_ps_lw).backward() if args.cd_ps_lw > 0: cd_inputs, cd_targets = cd_train_loader.next_batch() cd_inputs = cd_inputs.to(DEVICE) for k, v in cd_targets.items(): cd_targets[k] = v.to(DEVICE) pap_old = args.pap args.pap = False outputs = mgn(get_model_input(cd_inputs, cd_targets)) args.pap = pap_old cd_psl = (ps_loss(outputs[-3], cd_targets['ps_label']) + ps_loss(outputs[-2], cd_targets['ps_label']) + ps_loss(outputs[-1], cd_targets['ps_label'])) / 3. (cd_psl * args.cd_ps_lw).backward() optimizer_1.step() outputs = mgn(get_model_input(inputs, target)) optimizer_2.zero_grad() losses_2 = [ triplet_semihard_loss(output, labels, epoch) for output in outputs[2:5] ] loss_2 = sum(losses_2) / len(losses_2) psl = 0 if args.src_ps_lw > 0: psl = (ps_loss(outputs[-3], target['ps_label']) + ps_loss(outputs[-2], target['ps_label']) + ps_loss(outputs[-1], target['ps_label'])) / 3. (loss_2 + psl * args.src_ps_lw).backward() if args.cd_ps_lw > 0: cd_inputs, cd_targets = cd_train_loader.next_batch() cd_inputs = cd_inputs.to(DEVICE) for k, v in cd_targets.items(): cd_targets[k] = v.to(DEVICE) pap_old = args.pap args.pap = False outputs = mgn(get_model_input(cd_inputs, cd_targets)) args.pap = pap_old cd_psl = (ps_loss(outputs[-3], cd_targets['ps_label']) + ps_loss(outputs[-2], cd_targets['ps_label']) + ps_loss(outputs[-1], cd_targets['ps_label'])) / 3. (cd_psl * args.cd_ps_lw).backward() optimizer_2.step() running_loss_1 += loss_1.item() running_loss_2 += loss_2.item() running_loss = running_loss + (loss_1.item() + loss_2.item()) / 2.0 print('%d/%d - %d/%d - loss: %f - ps_loss: %f - cd_ps_loss: %f' % (epoch + 1, epochs, i, len(train_loader), (loss_1.item() + loss_2.item()) / 2, psl.item() if isinstance(psl, torch.Tensor) else 0, cd_psl.item() if args.cd_ps_lw > 0 else 0)) print('epoch: %d/%d - loss1: %f' % (epoch + 1, epochs, running_loss_1 / len(train_loader))) print('epoch: %d/%d - loss2: %f' % (epoch + 1, epochs, running_loss_2 / len(train_loader))) # if (epoch + 1) % 50 == 0: # model_weight_file = '{}/model_weight.pth'.format(exp_dir) # save_model(mgn, model_weight_file) # mgn.eval() # for name, test_loader in zip(testset_names, test_loaders): # test(test_loader[0], test_loader[1], test_loader[2], test_loader[3], trainset_name, name, epoch) model_weight_file = '{}/model_weight.pth'.format(exp_dir) save_model(mgn, model_weight_file) mgn.eval() for name, test_loader in zip(testset_names, test_loaders): test(test_loader[0], test_loader[1], test_loader[2], test_loader[3], trainset_name, name, epoch)
False) ids.extend(id_sam.tolist()) else: id_sam = np.random.choice(self.person_infos[pid], self.K, True) ids.extend(id_sam.tolist()) yield ids def __len__(self): return self.iter_num if __name__ == "__main__": from torch.utils.data import DataLoader from market1501 import Market1501 import cv2 ds = Market1501('./dataset/Market-1501-v15.09.15/bounding_box_train') sampler1 = BalancedSampler(ds, 2, 4) sampler2 = BalancedSampler(ds, 2, 4) dl1 = DataLoader(ds, batch_sampler=sampler1, num_workers=1) dl2 = DataLoader(ds, batch_sampler=sampler2, num_workers=1) for jj in range(2): for i, ((imgs1, lbs1, ids1), (imgs2, lbs2, ids2)) in enumerate(zip(dl1, dl2)): print(i) print(lbs1) print(lbs2) if i == 4: break
def run(): batch_size = 32 train_transform = transforms.Compose([ transforms.Resize(144, interpolation=3), transforms.RandomCrop((256, 128)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_transform = transforms.Compose([ transforms.Resize((288, 144), interpolation=3), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) test_flip_transform = transforms.Compose([ transforms.Resize((288, 144), interpolation=3), functional.hflip, transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) train_dataset = Market1501(root + '/bounding_box_train', transform=train_transform) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) query_dataset = Market1501(root + '/query', transform=test_transform) query_flip_dataset = Market1501(root + '/query', transform=test_flip_transform) query_loader = DataLoader(query_dataset, batch_size=batch_size, shuffle=False) query_flip_loader = DataLoader(query_flip_dataset, batch_size=batch_size, shuffle=False) test_dataset = Market1501(root + '/bounding_box_test', transform=test_transform) test_flip_dataset = Market1501(root + '/bounding_box_test', transform=test_flip_transform) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) test_flip_loader = DataLoader(test_flip_dataset, batch_size=batch_size, shuffle=False) ide = IDE(num_classes=len(train_dataset.unique_ids)).to(DEVICE) criterion = nn.CrossEntropyLoss() params = [ { 'params': ide.backbone.parameters(), 'lr': 0.01 }, { 'params': ide.classifier.parameters(), 'lr': 0.1 }, ] optimizer = optim.SGD(params, momentum=0.9, weight_decay=5e-4, nesterov=True) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1) epochs = 50 for epoch in range(epochs): ide.train() scheduler.step() running_loss = 0.0 for i, data in enumerate(train_loader): inputs, labels = data inputs, labels = inputs.to(DEVICE), labels.to(DEVICE) optimizer.zero_grad() outputs = ide(inputs) loss = criterion(outputs[1], labels) loss.backward() optimizer.step() running_loss += loss.item() print('%d/%d - %d/%d - loss: %f' % (epoch, epochs, i, len(train_loader), loss.item())) print('epoch: %d/%d - loss: %f' % (epoch, epochs, running_loss / len(train_loader))) if epoch % 10 == 9: ide.eval() query = np.concatenate([ ide(inputs.to(DEVICE))[0].detach().cpu().numpy() for inputs, _ in query_loader ]) query_flip = np.concatenate([ ide(inputs.to(DEVICE))[0].detach().cpu().numpy() for inputs, _ in query_flip_loader ]) test = np.concatenate([ ide(inputs.to(DEVICE))[0].detach().cpu().numpy() for inputs, _ in test_loader ]) test_flip = np.concatenate([ ide(inputs.to(DEVICE))[0].detach().cpu().numpy() for inputs, _ in test_flip_loader ]) # dist = cdist((query + query_flip) / 2., (test + test_flip) / 2.) dist = cdist(normalize(query + query_flip), normalize(test + test_flip)) r = cmc(dist, query_dataset.ids, test_dataset.ids, query_dataset.cameras, test_dataset.cameras, separate_camera_set=False, single_gallery_shot=False, first_match_break=True) m_ap = mean_ap(dist, query_dataset.ids, test_dataset.ids, query_dataset.cameras, test_dataset.cameras) print('epoch[%d]: mAP=%f, r@1=%f, r@3=%f, r@5=%f, r@10=%f' % (epoch + 1, m_ap, r[0], r[2], r[4], r[9]))