Example #1
0
def main():
    #GENERAL
    root = "/home/kuru/Desktop/veri-gms-master/"
    train_dir = '/home/kuru/Desktop/veri-gms-master/VeRispan/image_train/'
    source = {'veri'}
    target = {'veri'}
    workers = 2
    height = 320
    width = 320
    train_sampler = 'RandomSampler'

    #AUGMENTATION
    random_erase = True
    jitter = True
    aug = True

    #OPTIMIZATION
    opt = 'adam'
    lr = 0.0003
    weight_decay = 5e-4
    momentum = 0.9
    sgd_damp = 0.0
    nesterov = True
    warmup_factor = 0.01
    warmup_method = 'linear'

    #HYPERPARAMETER
    max_epoch = 80
    start = 0
    train_batch_size = 16
    test_batch_size = 50

    #SCHEDULER
    lr_scheduler = 'multi_step'
    stepsize = [30, 60]
    gamma = 0.1

    #LOSS
    margin = 0.3
    num_instances = 6
    lambda_tri = 1

    #MODEL
    arch = 'resnet101_ibn_a'
    no_pretrained = False

    #TEST SETTINGS
    load_weights = '/home/kuru/Desktop/veri-gms-master/IBN-Net_pytorch0.4.1/resnet101_ibn_a.pth'
    #load_weights = None
    start_eval = 0
    eval_freq = -1

    #MISC
    use_gpu = True
    use_amp = True
    print_freq = 50
    seed = 1
    resume = ''
    save_dir = '/home/kuru/Desktop/veri-gms-master/logapex/'
    gpu_id = 0
    vis_rank = True
    query_remove = True
    evaluate = False

    dataset_kwargs = {
        'source_names': source,
        'target_names': target,
        'root': root,
        'height': height,
        'width': width,
        'train_batch_size': train_batch_size,
        'test_batch_size': test_batch_size,
        'train_sampler': train_sampler,
        'random_erase': random_erase,
        'color_jitter': jitter,
        'color_aug': aug
    }
    transform_kwargs = {
        'height': height,
        'width': width,
        'random_erase': random_erase,
        'color_jitter': jitter,
        'color_aug': aug
    }

    optimizer_kwargs = {
        'optim': opt,
        'lr': lr,
        'weight_decay': weight_decay,
        'momentum': momentum,
        'sgd_dampening': sgd_damp,
        'sgd_nesterov': nesterov
    }

    lr_scheduler_kwargs = {
        'lr_scheduler': lr_scheduler,
        'stepsize': stepsize,
        'gamma': gamma
    }

    use_gpu = torch.cuda.is_available()
    log_name = 'log_test.txt' if evaluate else 'log_train.txt'
    sys.stdout = Logger(osp.join(save_dir, log_name))
    print('Currently using GPU ', gpu_id)
    cudnn.benchmark = True

    print('Initializing image data manager')
    dataset = init_imgreid_dataset(root='/home/kuru/Desktop/veri-gms-master/',
                                   name='veri')
    train = []
    num_train_pids = 0
    num_train_cams = 0

    for img_path, pid, camid in dataset.train:
        path = img_path[52:77]
        #print(path)
        folder = path[1:4]
        pid += num_train_pids
        camid += num_train_cams
        train.append((path, folder, pid, camid))

    num_train_pids += dataset.num_train_pids
    num_train_cams += dataset.num_train_cams

    pid = 0
    pidx = {}
    for img_path, pid, camid in dataset.train:
        path = img_path[52:77]

        folder = path[1:4]
        pidx[folder] = pid
        pid += 1

    path = '/home/kuru/Desktop/veri-gms-master/gms/'
    pkl = {}
    entries = os.listdir(path)
    for name in entries:
        f = open((path + name), 'rb')
        if name == 'featureMatrix.pkl':
            s = name[0:13]
        else:
            s = name[0:3]
        pkl[s] = pickle.load(f)
        f.close

    transform_t = train_transforms(**transform_kwargs)

    data_tfr = vd(
        pkl_file='index.pkl',
        dataset=train,
        root_dir='/home/kuru/Desktop/veri-gms-master/VeRi/image_train/',
        transform=transform_t)
    trainloader = DataLoader(data_tfr,
                             sampler=None,
                             batch_size=train_batch_size,
                             shuffle=True,
                             num_workers=workers,
                             pin_memory=False,
                             drop_last=True)

    #data_tfr = vd(pkl_file='index.pkl', dataset = train, root_dir=train_dir,transform=transforms.Compose([Rescale(64),RandomCrop(32),ToTensor()]))
    #dataloader = DataLoader(data_tfr, batch_size=batch_size, shuffle=False, num_workers=0)

    print('Initializing test data manager')
    dm = ImageDataManager(use_gpu, **dataset_kwargs)
    testloader_dict = dm.return_dataloaders()

    print('Initializing model: {}'.format(arch))
    model = models.init_model(name=arch,
                              num_classes=num_train_pids,
                              loss={'xent', 'htri'},
                              last_stride=1,
                              pretrained=not no_pretrained,
                              use_gpu=use_gpu)
    print('Model size: {:.3f} M'.format(count_num_param(model)))

    if load_weights is not None:
        print("weights loaded")
        load_pretrained_weights(model, load_weights)

    model = (model).cuda() if use_gpu else model

    #model = nn.DataParallel(model).cuda() if use_gpu else model
    optimizer = init_optimizer(model, **optimizer_kwargs)
    #optimizer = init_optimizer(model)

    model, optimizer = amp.initialize(model,
                                      optimizer,
                                      opt_level="O2",
                                      keep_batchnorm_fp32=True,
                                      loss_scale="dynamic")
    model = nn.DataParallel(model).cuda() if use_gpu else model
    scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs)

    criterion_xent = CrossEntropyLoss(num_classes=num_train_pids,
                                      use_gpu=use_gpu,
                                      label_smooth=True)
    criterion_htri = TripletLoss(margin=margin)
    ranking_loss = nn.MarginRankingLoss(margin=margin)

    if evaluate:
        print('Evaluate only')

        for name in target:
            print('Evaluating {} ...'.format(name))
            queryloader = testloader_dict[name]['query']
            galleryloader = testloader_dict[name]['gallery']
            _, distmat = test(model,
                              queryloader,
                              galleryloader,
                              train_batch_size,
                              use_gpu,
                              return_distmat=True)

            if vis_rank:
                visualize_ranked_results(distmat,
                                         dm.return_testdataset_by_name(name),
                                         save_dir=osp.join(
                                             save_dir, 'ranked_results', name),
                                         topk=20)
        return

    time_start = time.time()
    ranklogger = RankLogger(source, target)
    print('=> Start training')

    data_index = search(pkl)

    for epoch in range(start, max_epoch):
        losses = AverageMeter()
        #xent_losses = AverageMeter()
        htri_losses = AverageMeter()
        accs = AverageMeter()
        batch_time = AverageMeter()

        model.train()
        for p in model.parameters():
            p.requires_grad = True  # open all layers

        end = time.time()
        for batch_idx, (img, label, index, pid, cid) in enumerate(trainloader):
            trainX, trainY = torch.zeros(
                (train_batch_size * 3, 3, height, width),
                dtype=torch.float32), torch.zeros((train_batch_size * 3),
                                                  dtype=torch.int64)
            #pids = torch.zeros((batch_size*3), dtype = torch.int16)
            for i in range(train_batch_size):

                labelx = str(label[i])
                indexx = int(index[i])
                cidx = int(pid[i])
                if indexx > len(pkl[labelx]) - 1:
                    indexx = len(pkl[labelx]) - 1

                #maxx = np.argmax(pkl[labelx][indexx])
                a = pkl[labelx][indexx]
                minpos = np.argmax(ma.masked_where(a == 0, a))
                pos_dic = data_tfr[data_index[cidx][1] + minpos]

                neg_label = int(labelx)
                while True:
                    neg_label = random.choice(range(1, 770))
                    if neg_label is not int(labelx) and os.path.isdir(
                            os.path.join(
                                '/home/kuru/Desktop/adiusb/veri-split/train',
                                strint(neg_label))) is True:
                        break
                negative_label = strint(neg_label)
                neg_cid = pidx[negative_label]
                neg_index = random.choice(range(0, len(pkl[negative_label])))

                neg_dic = data_tfr[data_index[neg_cid][1] + neg_index]
                trainX[i] = img[i]
                trainX[i + train_batch_size] = pos_dic[0]
                trainX[i + (train_batch_size * 2)] = neg_dic[0]
                trainY[i] = cidx
                trainY[i + train_batch_size] = pos_dic[3]
                trainY[i + (train_batch_size * 2)] = neg_dic[3]

                #print("anc",labelx,'posdic', pos_dic[1],pos_dic[2],'negdic', neg_dic[1],neg_dic[2])

            trainX = trainX.cuda()
            trainY = trainY.cuda()
            outputs, features = model(trainX)
            xent_loss = criterion_xent(outputs[0:train_batch_size],
                                       trainY[0:train_batch_size])
            htri_loss = criterion_htri(features, trainY)

            #tri_loss = ranking_loss(features)
            #ent_loss = xent_loss(outputs[0:batch_size], trainY[0:batch_size], num_train_pids)

            loss = htri_loss + xent_loss
            optimizer.zero_grad()

            if use_amp:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

            #loss.backward()
            optimizer.step()

            batch_time.update(time.time() - end)
            losses.update(loss.item(), trainY.size(0))
            htri_losses.update(htri_loss.item(), trainY.size(0))
            accs.update(
                accuracy(outputs[0:train_batch_size],
                         trainY[0:train_batch_size])[0])

            if (batch_idx) % print_freq == 0:
                print('Train ', end=" ")
                print('Epoch: [{0}][{1}/{2}]\t'
                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                      'Acc {acc.val:.2f} ({acc.avg:.2f})\t'.format(
                          epoch + 1,
                          batch_idx + 1,
                          len(trainloader),
                          batch_time=batch_time,
                          loss=htri_losses,
                          acc=accs))

            end = time.time()

        scheduler.step()
        print('=> Test')

        for name in target:
            print('Evaluating {} ...'.format(name))
            queryloader = testloader_dict[name]['query']
            galleryloader = testloader_dict[name]['gallery']
            rank1, distmat = test(model, queryloader, galleryloader,
                                  test_batch_size, use_gpu)
            ranklogger.write(name, epoch + 1, rank1)
            rank2, distmat2 = test_rerank(model, queryloader, galleryloader,
                                          test_batch_size, use_gpu)
            ranklogger.write(name, epoch + 1, rank2)
        del queryloader
        del galleryloader
        del distmat
        #print(torch.cuda.memory_allocated(),torch.cuda.memory_cached())
        torch.cuda.empty_cache()

        if (epoch + 1) == max_epoch:
            #if (epoch + 1) % 10 == 0:
            print('=> Test')
            save_checkpoint(
                {
                    'state_dict': model.state_dict(),
                    'rank1': rank1,
                    'epoch': epoch + 1,
                    'arch': arch,
                    'optimizer': optimizer.state_dict(),
                }, save_dir)

            if vis_rank:
                visualize_ranked_results(distmat,
                                         dm.return_testdataset_by_name(name),
                                         save_dir=osp.join(
                                             save_dir, 'ranked_results', name),
                                         topk=20)
def main():
    #GENERAL
    torch.cuda.empty_cache()

    root = "/home/kuru/Desktop/veri-gms-master_noise/"
    train_dir = '/home/kuru/Desktop/veri-gms-master_noise/VeRispan/image_train/'
    source = {'verispan'}
    target = {'verispan'}
    workers = 4
    height = 280
    width  = 280
    train_size = 32
    train_sampler = 'RandomSampler'

    #AUGMENTATION
    random_erase = True
    jitter = True
    aug = True

    #OPTIMIZATION
    opt = 'adam'
    lr = 0.0003
    weight_decay = 5e-4
    momentum = 0.9
    sgd_damp = 0.0
    nesterov = True
    warmup_factor = 0.01
    warmup_method = 'linear'

    #HYPERPARAMETER
    max_epoch = 80
    start = 0
    train_batch_size = 8
    test_batch_size = 100

    #SCHEDULER
    lr_scheduler = 'multi_step'
    stepsize = [30, 60]
    gamma = 0.1

    #LOSS
    margin = 0.3
    num_instances = 4
    lambda_tri = 1

    #MODEL
    #arch = 'resnet101'
    arch='resnet101_ibn_a'
    no_pretrained = False

    #TEST SETTINGS
    load_weights = '/home/kuru/Desktop/veri-gms-master/IBN-Net_pytorch0.4.1/resnet101_ibn_a.pth'
    #load_weights = None
    start_eval = 0
    eval_freq = -1

    #MISC
    use_gpu = True
    print_freq = 10
    seed = 1
    resume = ''
    save_dir = '/home/kuru/Desktop/veri-gms-master_noise/spanningtree_verinoise_101_stride2/'
    gpu_id = 0,1
    vis_rank = True
    query_remove = True
    evaluate = False

    dataset_kwargs = {
        'source_names': source,
        'target_names': target,
        'root': root,
        'height': height,
        'width': width,
        'train_batch_size': train_batch_size,
        'test_batch_size': test_batch_size,
        'train_sampler': train_sampler,
        'random_erase': random_erase,
        'color_jitter': jitter,
        'color_aug': aug
        }
    transform_kwargs = {
        'height': height,
        'width': width,
        'random_erase': random_erase,
        'color_jitter': jitter,
        'color_aug': aug
    }

    optimizer_kwargs = {
        'optim': opt,
        'lr': lr,
        'weight_decay': weight_decay,
        'momentum': momentum,
        'sgd_dampening': sgd_damp,
        'sgd_nesterov': nesterov
        }

    lr_scheduler_kwargs = {
        'lr_scheduler': lr_scheduler,
        'stepsize': stepsize,
        'gamma': gamma
        }
    
    use_gpu = torch.cuda.is_available()
    log_name = 'log_test.txt' if evaluate else 'log_train.txt'
    sys.stdout = Logger(osp.join(save_dir, log_name))
    print('Currently using GPU ', gpu_id)
    cudnn.benchmark = True

    print('Initializing image data manager')
    dataset = init_imgreid_dataset(root='/home/kuru/Desktop/veri-gms-master_noise/', name='verispan')
    train = []
    num_train_pids = 0
    num_train_cams = 0

    print(len( dataset.train))

    for img_path, pid, camid, subid, countid in dataset.train:
        #print(img_path)
        path = img_path[56+6:90+6]
        #print(path)
        folder = path[1:4]
        #print(folder)
        pid += num_train_pids
        newidd=0
        train.append((path, folder, pid, camid,subid,countid))

    num_train_pids += dataset.num_train_pids
    num_train_cams += dataset.num_train_cams

    pid = 0
    pidx = {}
    for img_path, pid, camid, subid, countid in dataset.train:
        path = img_path[56+6:90+6]
        
        folder = path[1:4]
        pidx[folder] = pid
        pid+= 1

    sub=[]
    final=0
    xx=dataset.train
    newids=[]
    print(train[0:2])
    train2={}
    for k in range(0,770):
        for img_path, pid, camid, subid, countid in dataset.train:
            if k==pid:
                newid=final+subid
                sub.append(newid)
                #print(pid,subid,newid)
                newids.append(newid)
                train2[img_path]= newid
                #print(img_path, pid, camid, subid, countid, newid)

                

        final=max(sub)
        #print(final)
    print(len(newids),final)

    #train=train2
    #print(train2)
    train3=[]
    for img_path, pid, camid, subid, countid in dataset.train:
        #print(img_path,pid,train2[img_path])
        path = img_path[56:90+6]
        #print(path)
        folder = path[1:4]
        newid=train2[img_path]
        #print((path, folder, pid, camid, subid, countid,newid ))
        train3.append((path, folder, pid, camid, subid, countid,newid ))

    train = train3

    
    path = '/home/kuru/Desktop/adhi/veri-final-draft-master_noise/gmsNoise776/'
    pkl = {}
    #pkl[0] = pickle.load('/home/kuru/Desktop/veri-gms-master/gms/620.pkl')

    entries = os.listdir(path)
    for name in entries:
        f = open((path+name), 'rb')
        ccc=(path+name)
        #print(ccc)
        if name=='featureMatrix.pkl':
            s = name[0:13]
        else:
            s = name[0:3]
        #print(s)
        #with open (ccc,"rb") as ff:
        #    pkl[s] = pickle.load(ff)
            #print(pkl[s])
        pkl[s] = pickle.load(f)
        f.close
        #print(len(pkl))

    with open('cids.pkl', 'rb') as handle:
        b = pickle.load(handle)
        #print(b)

    with open('index.pkl', 'rb') as handle:
        c = pickle.load(handle)



    transform_t = train_transforms(**transform_kwargs)

    data_tfr = vdspan(pkl_file='index_veryspan_noise.pkl', dataset = train, root_dir='/home/kuru/Desktop/veri-gms-master_noise/VeRispan/image_train/', transform=transform_t)
    print("lllllllllllllllllllllllllllllllllllllllllllline 433")
    df2=[]
    data_tfr_old=data_tfr
    for (img,label,index,pid, cid,subid,countid,newid) in data_tfr :
        #print((img,label,index,pid, cid,subid,countid,newid) )
        #print("datframe",(label))
        #print(countid)
        if countid > 4 :
            #print(countid)
            df2.append((img,label,index,pid, cid,subid,countid,newid))
    print("filtered final trainset length",len(df2))
    
    data_tfr=df2
    
    
    
    
    trainloader = DataLoader(data_tfr, sampler=None,batch_size=train_batch_size, shuffle=True, num_workers=workers,pin_memory=True, drop_last=True)

    #data_tfr = vd(pkl_file='index.pkl', dataset = train, root_dir=train_dir,transform=transforms.Compose([Rescale(64),RandomCrop(32),ToTensor()]))
    #dataloader = DataLoader(data_tfr, batch_size=batch_size, shuffle=False, num_workers=0)

    for batch_idx, (img,label,index,pid, cid,subid,countid,newid) in enumerate(trainloader):
        #print("trainloader",batch_idx, (label,index,pid, cid,subid,countid,newid))
        print("trainloader",batch_idx, (label))
        break

    print('Initializing test data manager')
    dm = ImageDataManager(use_gpu, **dataset_kwargs)
    testloader_dict = dm.return_dataloaders()

    print('Initializing model: {}'.format(arch))
    model = models.init_model(name=arch, num_classes=num_train_pids, loss={'xent', 'htri'},
                              pretrained=not no_pretrained, last_stride =2 )
    print('Model size: {:.3f} M'.format(count_num_param(model)))

    if load_weights is not None:
        print("weights loaded")
        load_pretrained_weights(model, load_weights)

    print(torch.cuda.device_count())
    model = nn.DataParallel(model).cuda() if use_gpu else model
    optimizer = init_optimizer(model, **optimizer_kwargs)
    #optimizer = init_optimizer(model)
    
    scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs)

    criterion_xent = CrossEntropyLoss(num_classes=num_train_pids, use_gpu=use_gpu, label_smooth=True)
    criterion_htri = TripletLoss(margin=margin)
    ranking_loss = nn.MarginRankingLoss(margin = margin)

    if evaluate:
        print('Evaluate only')

        for name in target:
            print('Evaluating {} ...'.format(name))
            queryloader = testloader_dict[name]['query']
            galleryloader = testloader_dict[name]['gallery']
            _, distmat = test(model, queryloader, galleryloader, train_batch_size, use_gpu, return_distmat=True)

            if vis_rank:
                visualize_ranked_results(
                    distmat, dm.return_testdataset_by_name(name),
                    save_dir=osp.join(save_dir, 'ranked_results', name),
                    topk=20
                )
        return    

    time_start = time.time()
    ranklogger = RankLogger(source, target)
    print('=> Start training')

    data_index = search(pkl)
    print(len(data_index))
    
    for epoch in range(start, max_epoch):
        losses = AverageMeter()
        #xent_losses = AverageMeter()
        htri_losses = AverageMeter()
        accs = AverageMeter()
        batch_time = AverageMeter()
        xent_losses=AverageMeter()

        model.train()
        for p in model.parameters():
            p.requires_grad = True    # open all layers

        end = time.time()
        for batch_idx,  (img,label,index,pid, cid,subid,countid,newid)  in enumerate(trainloader):
            trainX, trainY = torch.zeros((train_batch_size*3,3,height, width), dtype=torch.float32), torch.zeros((train_batch_size*3), dtype = torch.int64)
            #pids = torch.zeros((batch_size*3), dtype = torch.int16)
            for i in range(train_batch_size):
                #print("dfdsfs")
                labelx = label[i]
                indexx = index[i]
                cidx = pid[i]
                if indexx >len(pkl[labelx])-1:
                    indexx = len(pkl[labelx])-1

                #maxx = np.argmax(pkl[labelx][indexx])
                a = pkl[labelx][indexx]
                minpos = np.argmin(ma.masked_where(a==0, a)) 
                #print(minpos)
                #print(np.array(data_index).shape)
                #print(data_index[cidx][1])
                pos_dic = data_tfr_old[data_index[cidx][1]+minpos]

                neg_label = int(labelx)
                while True:
                    neg_label = random.choice(range(1, 770))
                    #print(neg_label)
                    if neg_label is not int(labelx) and os.path.isdir(os.path.join('/home/kuru/Desktop/adiusb/veri-split/train', strint(neg_label))) is True:
                        break
                negative_label = strint(neg_label)
                neg_cid = pidx[negative_label]
                neg_index = random.choice(range(0, len(pkl[negative_label])))

                neg_dic = data_tfr_old[data_index[neg_cid][1]+neg_index]
                trainX[i] = img[i]
                trainX[i+train_batch_size] = pos_dic[0]
                trainX[i+(train_batch_size*2)] = neg_dic[0]
                trainY[i] = cidx
                trainY[i+train_batch_size] = pos_dic[3]
                trainY[i+(train_batch_size*2)] = neg_dic[3]
            
            trainX = trainX.cuda()
            trainY = trainY.cuda()
            outputs, features = model(trainX)
            xent_loss = criterion_xent(outputs[0:train_batch_size], trainY[0:train_batch_size])
            htri_loss = criterion_htri(features, trainY)

            #tri_loss = ranking_loss(features)
            #ent_loss = xent_loss(outputs[0:batch_size], trainY[0:batch_size], num_train_pids)
            
            loss = htri_loss+xent_loss
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            batch_time.update(time.time() - end)
            losses.update(loss.item(), trainY.size(0))
            htri_losses.update(htri_loss.item(), trainY.size(0))
            xent_losses.update(xent_loss.item(), trainY.size(0))
            accs.update(accuracy(outputs[0:train_batch_size], trainY[0:train_batch_size])[0])
    
            if (batch_idx) % 50 == 0:
                print('Train ', end=" ")
                print('Epoch: [{0}][{1}/{2}]\t'
                    'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                    'TriLoss {loss.val:.4f} ({loss.avg:.4f})\t'
                    'XLoss {xloss.val:.4f} ({xloss.avg:.4f})\t'
                    'OveralLoss {oloss.val:.4f} ({oloss.avg:.4f})\t'
                    'Acc {acc.val:.2f} ({acc.avg:.2f})\t'
                    'lr {lrrr} \t'.format(
                    epoch + 1, batch_idx + 1, len(trainloader),
                    batch_time=batch_time,
                    loss = htri_losses,
                    xloss = xent_losses,
                    oloss = losses,
                    acc=accs ,
                    lrrr=lrrr,
                ))
                

            end = time.time()

        
        scheduler.step()            
        print('=> Test')

        for name in target:
            print('Evaluating {} ...'.format(name))
            queryloader = testloader_dict[name]['query']
            galleryloader = testloader_dict[name]['gallery']
            rank1, distmat = test(model, queryloader, galleryloader, test_batch_size, use_gpu)
            ranklogger.write(name, epoch + 1, rank1)
            rank2, distmat2 = test_rerank(model, queryloader, galleryloader, test_batch_size, use_gpu)
            ranklogger.write(name, epoch + 1, rank2)
            
        #if (epoch + 1) == max_epoch:
        if (epoch + 1) % 2 == 0:
            print('=> Test')

            for name in target:
                print('Evaluating {} ...'.format(name))
                queryloader = testloader_dict[name]['query']
                galleryloader = testloader_dict[name]['gallery']
                rank1, distmat = test(model, queryloader, galleryloader, test_batch_size, use_gpu)
                ranklogger.write(name, epoch + 1, rank1)

                # if vis_rank:
                #     visualize_ranked_results(
                #         distmat, dm.return_testdataset_by_name(name),
                #         save_dir=osp.join(save_dir, 'ranked_results', name),
                #         topk=20)

            save_checkpoint({
                'state_dict': model.state_dict(),
                'rank1': rank1,
                'epoch': epoch + 1,
                'arch': arch,
                'optimizer': optimizer.state_dict(),
            }, save_dir)
Example #3
0
def main():
    #GENERAL
    torch.cuda.empty_cache()
    root = "/home/kuru/Desktop/veri-gms-master/"
    train_dir = '/home/kuru/Desktop/veri-gms-master/VeRispan/image_train/'
    source = {'verispan'}
    target = {'verispan'}
    workers = 4
    height = 320
    width = 320
    train_sampler = 'RandomSampler'

    #AUGMENTATION
    random_erase = True
    jitter = True
    aug = True

    #OPTIMIZATION
    opt = 'adam'
    lr = 0.001
    weight_decay = 5e-4
    momentum = 0.9
    sgd_damp = 0.0
    nesterov = True
    warmup_factor = 0.01
    warmup_method = 'linear'

    STEPS = (30, 60)
    GAMMA = 0.1
    WARMUP_FACTOR = 0.01
    WARMUP_EPOCHS = 10
    WARMUP_METHOD = 'linear'

    #HYPERPARAMETER
    max_epoch = 80
    start = 0
    train_batch_size = 16
    test_batch_size = 50

    #SCHEDULER
    lr_scheduler = 'multi_step'
    stepsize = [30, 60]
    gamma = 0.1

    #LOSS
    margin = 0.3
    num_instances = 4
    lambda_tri = 1

    #MODEL
    #arch = 'resnet101'
    arch = 'resnet50_ibn_a'
    no_pretrained = False

    #TEST SETTINGS
    #load_weights = '/home/kuru/Desktop/veri-gms-master/IBN-Net_pytorch0.4.1/resnet101_ibn_a.pth'

    #load_weights = '/home/kuru/Desktop/veri-gms-master/IBN-Net_pytorch0.4.1/resnet101_ibn_a.pth'
    load_weights = '/home/kuru/Desktop/veri-gms-master/IBN-Net_pytorch0.4.1/resnet50_ibn_a.pth'

    #load_weights = None
    start_eval = 0
    eval_freq = -1

    num_classes = 776
    feat_dim = 2048
    CENTER_LR = 0.5
    CENTER_LOSS_WEIGHT = 0.0005
    center_criterion = CenterLoss(num_classes=num_classes,
                                  feat_dim=feat_dim,
                                  use_gpu=True)
    optimizer_center = torch.optim.SGD(center_criterion.parameters(),
                                       lr=CENTER_LR)

    #MISC
    use_gpu = True
    #use_gpu = False
    print_freq = 10
    seed = 1
    resume = ''
    save_dir = '/home/kuru/Desktop/veri-gms-master_noise/spanningtree_veri_pure/'
    gpu_id = 0, 1
    vis_rank = True
    query_remove = True
    evaluate = False

    dataset_kwargs = {
        'source_names': source,
        'target_names': target,
        'root': root,
        'height': height,
        'width': width,
        'train_batch_size': train_batch_size,
        'test_batch_size': test_batch_size,
        'train_sampler': train_sampler,
        'random_erase': random_erase,
        'color_jitter': jitter,
        'color_aug': aug
    }
    transform_kwargs = {
        'height': height,
        'width': width,
        'random_erase': random_erase,
        'color_jitter': jitter,
        'color_aug': aug
    }

    optimizer_kwargs = {
        'optim': opt,
        'lr': lr,
        'weight_decay': weight_decay,
        'momentum': momentum,
        'sgd_dampening': sgd_damp,
        'sgd_nesterov': nesterov
    }

    lr_scheduler_kwargs = {
        'lr_scheduler': lr_scheduler,
        'stepsize': stepsize,
        'gamma': gamma
    }

    use_gpu = torch.cuda.is_available()

    log_name = 'log_test.txt' if evaluate else 'log_train.txt'
    sys.stdout = Logger(osp.join(save_dir, log_name))
    print('Currently using GPU ', gpu_id)
    cudnn.benchmark = True

    print('Initializing image data manager')
    #dataset = init_imgreid_dataset(root='/home/kuru/Desktop/veri-gms-master/', name='veri')
    dataset = init_imgreid_dataset(root='/home/kuru/Desktop/veri-gms-master/',
                                   name='verispan')
    train = []
    num_train_pids = 0
    num_train_cams = 0
    print(len(dataset.train))

    for img_path, pid, camid, subid, countid in dataset.train:
        #print(img_path)
        path = img_path[56:90 + 6]
        #print(path)
        folder = path[1:4]
        #print(folder)
        #print(img_path, pid, camid,subid,countid)
        pid += num_train_pids
        camid += num_train_cams
        newidd = 0
        train.append((path, folder, pid, camid, subid, countid))
        #print(train)
        #break

    num_train_pids += dataset.num_train_pids
    num_train_cams += dataset.num_train_cams

    pid = 0
    pidx = {}
    for img_path, pid, camid, subid, countid in dataset.train:
        path = img_path[56:90 + 6]

        folder = path[1:4]
        pidx[folder] = pid
        pid += 1
    #print(pidx)

    sub = []
    final = 0
    xx = dataset.train
    newids = []
    print(train[0:2])
    train2 = {}
    for k in range(0, 770):
        for img_path, pid, camid, subid, countid in dataset.train:
            if k == pid:
                newid = final + subid
                sub.append(newid)
                #print(pid,subid,newid)
                newids.append(newid)
                train2[img_path] = newid
                #print(img_path, pid, camid, subid, countid, newid)

        final = max(sub)
        #print(final)
    print(len(newids), final)

    #train=train2
    #print(train2)
    train3 = []
    for img_path, pid, camid, subid, countid in dataset.train:
        #print(img_path,pid,train2[img_path])
        path = img_path[56:90 + 6]
        #print(path)
        folder = path[1:4]
        newid = train2[img_path]
        #print((path, folder, pid, camid, subid, countid,newid ))
        train3.append((path, folder, pid, camid, subid, countid, newid))

    train = train3

    # for (path, folder, pid, camid, subid, countid,newid) in train:
    #     print(path, folder)

    #path = '/home/kuru/Desktop/adhi/veri-final-draft-master_noise/gmsNoise776/'
    path = '/home/kuru/Desktop/veri-gms-master/gms/'
    pkl = {}
    #pkl[0] = pickle.load('/home/kuru/Desktop/veri-gms-master/gms/620.pkl')

    entries = os.listdir(path)
    for name in entries:
        f = open((path + name), 'rb')
        ccc = (path + name)
        #print(ccc)
        if name == 'featureMatrix.pkl':
            s = name[0:13]
        else:
            s = name[0:3]
        #print(s)
        #with open (ccc,"rb") as ff:
        #    pkl[s] = pickle.load(ff)
        #print(pkl[s])
        pkl[s] = pickle.load(f)
        f.close
        #print(len(pkl))

    print('=> pickle indexing')

    data_index = search(pkl)
    print(len(data_index))

    transform_t = train_transforms(**transform_kwargs)
    #print(train[0],train[10])

    #data_tfr = vd(pkl_file='index.pkl', dataset = train, root_dir='/home/kuru/Desktop/veri-gms-master/VeRi/image_train/', transform=transform_t)
    data_tfr = vdspan(
        pkl_file='index_veryspan.pkl',
        dataset=train,
        root_dir='/home/kuru/Desktop/veri-gms-master/VeRispan/image_train/',
        transform=transform_t)
    #print(data_tfr)
    #print(trainloader)
    #data_tfr2=list(data_tfr)
    print("lllllllllllllllllllllllllllllllllllllllllllline 433")
    df2 = []
    data_tfr_old = data_tfr
    for (img, label, index, pid, cid, subid, countid, newid) in data_tfr:
        #print((img,label,index,pid, cid,subid,countid,newid) )
        #print("datframe",(label))
        #print(countid)
        if countid > 4:
            #print(countid)
            df2.append((img, label, index, pid, cid, subid, countid, newid))
    print("filtered final trainset length", len(df2))

    data_tfr = df2

    # with open('df2noise_ex.pkl', 'wb') as handle:
    #     b = pickle.dump(df2, handle, protocol=pickle.HIGHEST_PROTOCOL)

    # with open('df2noise.pkl', 'rb') as handle:
    #     df2 = pickle.load(handle)
    # data_tfr=df2
    # for (img,label,index,pid, cid,subid,countid,newid) in data_tfr :
    #     print("datframe",(label))

    #data_tfr = vdspansort( dataset = train, root_dir='/home/kuru/Desktop/veri-gms-master_noise/VeRispan/image_train/', transform=transform_t)

    #trainloader = DataLoader(df2, sampler=None,batch_size=train_batch_size, shuffle=True, num_workers=workers,pin_memory=True, drop_last=True)
    trainloader = DataLoader(data_tfr,
                             sampler=None,
                             batch_size=train_batch_size,
                             shuffle=True,
                             num_workers=workers,
                             pin_memory=True,
                             drop_last=True)

    for batch_idx, (img, label, index, pid, cid, subid, countid,
                    newid) in enumerate(trainloader):
        #print("trainloader",batch_idx, (label,index,pid, cid,subid,countid,newid))
        print("trainloader", batch_idx, (label))
        break

    print('Initializing test data manager')
    dm = ImageDataManager(use_gpu, **dataset_kwargs)
    testloader_dict = dm.return_dataloaders()

    print('Initializing model: {}'.format(arch))
    model = models.init_model(name=arch,
                              num_classes=num_train_pids,
                              loss={'xent', 'htri'},
                              pretrained=not no_pretrained,
                              last_stride=2)
    print('Model size: {:.3f} M'.format(count_num_param(model)))

    if load_weights is not None:
        print("weights loaded")
        load_pretrained_weights(model, load_weights)

    #checkpoint = torch.load('/home/kuru/Desktop/veri-gms-master/logg/model.pth.tar-19')
    #model._load_from_state_dict(checkpoint['state_dict'])
    #model.load_state_dict(checkpoint['state_dict'])

    #optimizer.load_state_dict(checkpoint['optimizer'])
    #print(checkpoint['epoch'])
    #print(checkpoint['rank1'])
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    print(torch.cuda.device_count())
    model = nn.DataParallel(model).cuda() if use_gpu else model
    optimizer = init_optimizer(model, **optimizer_kwargs)

    #optimizer = init_optimizer(model)
    #optimizer.load_state_dict(checkpoint['optimizer'])

    scheduler = init_lr_scheduler(optimizer, **lr_scheduler_kwargs)
    # scheduler = WarmupMultiStepLR(optimizer, STEPS, GAMMA,
    #                               WARMUP_FACTOR,
    #                               WARMUP_EPOCHS, WARMUP_METHOD)

    criterion_xent = CrossEntropyLoss(num_classes=num_train_pids,
                                      use_gpu=use_gpu,
                                      label_smooth=True)
    criterion_htri = TripletLoss(margin=margin)
    ranking_loss = nn.MarginRankingLoss(margin=margin)

    if evaluate:
        print('Evaluate only')

        for name in target:
            print('Evaluating {} ...'.format(name))
            queryloader = testloader_dict[name]['query']
            galleryloader = testloader_dict[name]['gallery']
            _, distmat = test(model,
                              queryloader,
                              galleryloader,
                              train_batch_size,
                              use_gpu,
                              return_distmat=True)

            if vis_rank:
                visualize_ranked_results(distmat,
                                         dm.return_testdataset_by_name(name),
                                         save_dir=osp.join(
                                             save_dir, 'ranked_results', name),
                                         topk=20)
        return

    time_start = time.time()
    ranklogger = RankLogger(source, target)

    # # checkpoint = torch.load('/home/kuru/Desktop/market_all/ibna_model/model.pth.tar-79')
    # # model.load_state_dict(checkpoint['state_dict'])
    # # optimizer.load_state_dict(checkpoint['optimizer'])
    # # print(checkpoint['epoch'])
    # # start_epoch=checkpoint['epoch']
    # # start=start_epoch

    # checkpoint = torch.load('/home/kuru/Desktop/veri-gms-master/spanningtreeveri/model.pth.tar-2')
    # model.load_state_dict(checkpoint['state_dict'])
    # optimizer.load_state_dict(checkpoint['optimizer'])
    # print(checkpoint['epoch'])
    # start_epoch=checkpoint['epoch']
    # start=start_epoch

    ##start_epoch=resume_from_checkpoint('/home/kuru/Desktop/veri-gms-master/logg/model.pth.tar-20', model, optimizer=None)
    print('=> Start training')

    for epoch in range(start, max_epoch):
        print(epoch, scheduler.get_lr()[0])
        #print( torch.cuda.memory_allocated(0))
        losses = AverageMeter()
        #xent_losses = AverageMeter()
        htri_losses = AverageMeter()
        accs = AverageMeter()
        batch_time = AverageMeter()
        xent_losses = AverageMeter()

        model.train()
        for p in model.parameters():
            p.requires_grad = True  # open all layers

        end = time.time()
        for batch_idx, (img, label, index, pid, cid, subid, countid,
                        newid) in enumerate(trainloader):
            trainX, trainY = torch.zeros(
                (train_batch_size * 3, 3, height, width),
                dtype=torch.float32), torch.zeros((train_batch_size * 3),
                                                  dtype=torch.int64)
            #pids = torch.zeros((batch_size*3), dtype = torch.int16)
            #batchcount=0
            for i in range(train_batch_size):
                if (countid[i] > 4):
                    #batchcount=batchcount+1
                    #print("dfdsfs")
                    labelx = label[i]
                    indexx = index[i]
                    cidx = pid[i]
                    if indexx > len(pkl[labelx]) - 1:
                        indexx = len(pkl[labelx]) - 1

                    #maxx = np.argmax(pkl[labelx][indexx])
                    a = pkl[labelx][indexx]
                    minpos = np.argmin(ma.masked_where(a == 0, a))

                    # print(len(a))
                    # print(a)
                    # print(ma.masked_where(a==0, a))
                    # print(labelx,index,pid,cidx,minpos)
                    # print(np.array(data_index).shape)
                    # print(data_index[cidx][1])
                    pos_dic = data_tfr_old[data_index[cidx][1] + minpos]
                    #print('posdic', pos_dic)

                    neg_label = int(labelx)
                    while True:
                        neg_label = random.choice(range(1, 770))
                        #print(neg_label)
                        if neg_label is not int(labelx) and os.path.isdir(
                                os.path.join(
                                    '/home/kuru/Desktop/veri-gms-master_noise/veriNoise_train_spanning_folder',
                                    strint(neg_label))) is True:
                            break
                    negative_label = strint(neg_label)
                    neg_cid = pidx[negative_label]
                    neg_index = random.choice(
                        range(0, len(pkl[negative_label])))
                    #print(negative_label,neg_cid,neg_index,data_index[neg_cid] )
                    neg_dic = data_tfr_old[data_index[neg_cid][1] + neg_index]
                    #print('negdic', neg_dic)
                    trainX[i] = img[i]
                    trainX[i + train_batch_size] = pos_dic[0]
                    trainX[i + (train_batch_size * 2)] = neg_dic[0]
                    trainY[i] = cidx
                    trainY[i + train_batch_size] = pos_dic[3]
                    trainY[i + (train_batch_size * 2)] = neg_dic[3]
                    # trainY[i+train_batch_size] = pos_dic[7]
                    # trainY[i+(train_batch_size*2)] = neg_dic[7]
                #break
                # else:
                #     print("skiped",countid[i],subid[i],label[i])
            #break
            #print(batchcount)
            trainX = trainX.cuda()
            trainY = trainY.cuda()
            outputs, features = model(trainX)
            xent_loss = criterion_xent(outputs[0:train_batch_size],
                                       trainY[0:train_batch_size])
            htri_loss = criterion_htri(features, trainY)
            centerloss = CENTER_LOSS_WEIGHT * center_criterion(
                features, trainY)

            #tri_loss = ranking_loss(features)
            #ent_loss = xent_loss(outputs[0:batch_size], trainY[0:batch_size], num_train_pids)

            loss = htri_loss + xent_loss + centerloss
            loss = htri_loss + xent_loss

            optimizer.zero_grad()
            optimizer_center.zero_grad()
            loss.backward()
            optimizer.step()
            # for param in center_criterion.parameters():
            #     param.grad.data *= (1. /CENTER_LOSS_WEIGHT)
            # optimizer_center.step()

            for param_group in optimizer.param_groups:
                #print(param_group['lr'] )
                lrrr = str(param_group['lr'])

            batch_time.update(time.time() - end)
            losses.update(loss.item(), trainY.size(0))
            htri_losses.update(htri_loss.item(), trainY.size(0))
            xent_losses.update(xent_loss.item(), trainY.size(0))
            accs.update(
                accuracy(outputs[0:train_batch_size],
                         trainY[0:train_batch_size])[0])

            if (batch_idx) % 50 == 0:
                print('Train ', end=" ")
                print('Epoch: [{0}][{1}/{2}]\t'
                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'TriLoss {loss.val:.4f} ({loss.avg:.4f})\t'
                      'XLoss {xloss.val:.4f} ({xloss.avg:.4f})\t'
                      'OveralLoss {oloss.val:.4f} ({oloss.avg:.4f})\t'
                      'Acc {acc.val:.2f} ({acc.avg:.2f})\t'
                      'lr {lrrr} \t'.format(
                          epoch + 1,
                          batch_idx + 1,
                          len(trainloader),
                          batch_time=batch_time,
                          loss=htri_losses,
                          xloss=xent_losses,
                          oloss=losses,
                          acc=accs,
                          lrrr=lrrr,
                      ))

            end = time.time()

        # del loss
        # del htri_loss
        # del xent_loss
        # del htri_losses
        # del losses
        # del outputs
        # del features
        # del accs
        # del trainX
        # del trainY

        scheduler.step()
        print('=> Test')
        save_checkpoint(
            {
                'state_dict': model.state_dict(),
                #'rank1': rank1,
                'epoch': epoch + 1,
                'arch': arch,
                'optimizer': optimizer.state_dict(),
            },
            save_dir)
        GPUtil.showUtilization()
        print(torch.cuda.memory_allocated(), torch.cuda.memory_cached())
        for name in target:
            print('Evaluating {} ...'.format(name))
            queryloader = testloader_dict[name]['query']
            galleryloader = testloader_dict[name]['gallery']
            rank1, distmat = test(model, queryloader, galleryloader,
                                  test_batch_size, use_gpu)
            ranklogger.write(name, epoch + 1, rank1)
            rank2, distmat2 = test_rerank(model, queryloader, galleryloader,
                                          test_batch_size, use_gpu)
            ranklogger.write(name, epoch + 1, rank2)
        del queryloader
        del galleryloader
        del distmat
        print(torch.cuda.memory_allocated(), torch.cuda.memory_cached())
        torch.cuda.empty_cache()

        if (epoch + 1) == max_epoch:
            #if (epoch + 1) % 10 == 0:
            print('=> Test')
            save_checkpoint(
                {
                    'state_dict': model.state_dict(),
                    'rank1': rank1,
                    'epoch': epoch + 1,
                    'arch': arch,
                    'optimizer': optimizer.state_dict(),
                }, save_dir)
            for name in target:
                print('Evaluating {} ...'.format(name))
                queryloader = testloader_dict[name]['query']
                galleryloader = testloader_dict[name]['gallery']
                rank1, distmat = test(model, queryloader, galleryloader,
                                      test_batch_size, use_gpu)
                ranklogger.write(name, epoch + 1, rank1)
                # del queryloader
                # del galleryloader
                # del distmat

                if vis_rank:
                    visualize_ranked_results(
                        distmat,
                        dm.return_testdataset_by_name(name),
                        save_dir=osp.join(save_dir, 'ranked_results', name),
                        topk=20)
Example #4
0
def main():

    torch.backends.cudnn.deterministic = True
    cudnn.benchmark = True
    #parser = argparse.ArgumentParser(description="ReID Baseline Training")
    #parser.add_argument(
    #"--config_file", default="", help="path to config file", type=str)

    #parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER)

    #args = parser.parse_args()
    config_file = 'configs/baseline_veri_r101_a.yml'
    if config_file != "":
        cfg.merge_from_file(config_file)
    #cfg.merge_from_list(args.opts)
    cfg.freeze()

    output_dir = cfg.OUTPUT_DIR
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    logger = setup_logger("reid_baseline", output_dir, if_train=True)
    logger.info("Saving model in the path :{}".format(cfg.OUTPUT_DIR))
    logger.info(config_file)

    if config_file != "":
        logger.info("Loaded configuration file {}".format(config_file))
        with open(config_file, 'r') as cf:
            config_str = "\n" + cf.read()
            logger.info(config_str)
    logger.info("Running with config:\n{}".format(cfg))

    os.environ['CUDA_VISIBLE_DEVICES'] = cfg.MODEL.DEVICE_ID

    path = 'D:/Python_SMU/Veri/verigms/gms/'
    pkl = {}
    entries = os.listdir(path)
    for name in entries:
        f = open((path + name), 'rb')
        if name == 'featureMatrix.pkl':
            s = name[0:13]
        else:
            s = name[0:3]
        pkl[s] = pickle.load(f)
        f.close

    with open('cids.pkl', 'rb') as handle:
        b = pickle.load(handle)

    with open('index.pkl', 'rb') as handle:
        c = pickle.load(handle)

    train_transforms, val_transforms, dataset, train_set, val_set = make_dataset(
        cfg, pkl_file='index.pkl')

    num_workers = cfg.DATALOADER.NUM_WORKERS
    num_classes = dataset.num_train_pids
    #pkl_f = 'index.pkl'
    pid = 0
    pidx = {}
    for img_path, pid, _, _ in dataset.train:
        path = img_path.split('\\')[-1]
        folder = path[1:4]
        pidx[folder] = pid
        pid += 1

    if 'triplet' in cfg.DATALOADER.SAMPLER:
        train_loader = DataLoader(train_set,
                                  batch_size=cfg.SOLVER.IMS_PER_BATCH,
                                  sampler=RandomIdentitySampler(
                                      dataset.train, cfg.SOLVER.IMS_PER_BATCH,
                                      cfg.DATALOADER.NUM_INSTANCE),
                                  num_workers=num_workers,
                                  pin_memory=True,
                                  collate_fn=train_collate_fn)
    elif cfg.DATALOADER.SAMPLER == 'softmax':
        print('using softmax sampler')
        train_loader = DataLoader(train_set,
                                  batch_size=cfg.SOLVER.IMS_PER_BATCH,
                                  shuffle=True,
                                  num_workers=num_workers,
                                  pin_memory=True,
                                  collate_fn=train_collate_fn)
    else:
        print('unsupported sampler! expected softmax or triplet but got {}'.
              format(cfg.SAMPLER))

    print("train loader loaded successfully")

    val_loader = DataLoader(val_set,
                            batch_size=cfg.TEST.IMS_PER_BATCH,
                            shuffle=False,
                            num_workers=num_workers,
                            pin_memory=True,
                            collate_fn=train_collate_fn)
    print("val loader loaded successfully")

    if cfg.MODEL.PRETRAIN_CHOICE == 'finetune':
        model = make_model(cfg, num_class=576)
        model.load_param_finetune(cfg.MODEL.PRETRAIN_PATH)
        print('Loading pretrained model for finetuning......')
    else:
        model = make_model(cfg, num_class=num_classes)

    loss_func, center_criterion = make_loss(cfg, num_classes=num_classes)

    optimizer, optimizer_center = make_optimizer(cfg, model, center_criterion)
    scheduler = WarmupMultiStepLR(optimizer, cfg.SOLVER.STEPS,
                                  cfg.SOLVER.GAMMA, cfg.SOLVER.WARMUP_FACTOR,
                                  cfg.SOLVER.WARMUP_EPOCHS,
                                  cfg.SOLVER.WARMUP_METHOD)

    print("model,optimizer, loss, scheduler loaded successfully")

    height, width = cfg.INPUT.SIZE_TRAIN

    log_period = cfg.SOLVER.LOG_PERIOD
    checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
    eval_period = cfg.SOLVER.EVAL_PERIOD

    device = "cuda"
    epochs = cfg.SOLVER.MAX_EPOCHS

    logger = logging.getLogger("reid_baseline.train")
    logger.info('start training')

    if device:
        if torch.cuda.device_count() > 1:
            print('Using {} GPUs for training'.format(
                torch.cuda.device_count()))
            model = nn.DataParallel(model)
        model.to(device)

    loss_meter = AverageMeter()
    acc_meter = AverageMeter()

    evaluator = R1_mAP_eval(len(dataset.query),
                            max_rank=50,
                            feat_norm=cfg.TEST.FEAT_NORM)
    model.base._freeze_stages()
    logger.info('Freezing the stages number:{}'.format(cfg.MODEL.FROZEN))

    data_index = search(pkl)
    print("Ready for training")

    for epoch in range(1, epochs + 1):
        start_time = time.time()
        loss_meter.reset()
        acc_meter.reset()
        evaluator.reset()
        scheduler.step()
        model.train()
        for n_iter, (img, label, index, pid, cid) in enumerate(train_loader):
            optimizer.zero_grad()
            optimizer_center.zero_grad()
            #img = img.to(device)
            #target = vid.to(device)
            trainX, trainY = torch.zeros(
                (train_loader.batch_size * 3, 3, height, width),
                dtype=torch.float32), torch.zeros(
                    (train_loader.batch_size * 3), dtype=torch.int64)

            for i in range(train_loader.batch_size):
                labelx = label[i]
                indexx = index[i]
                cidx = pid[i]
                if indexx > len(pkl[labelx]) - 1:
                    indexx = len(pkl[labelx]) - 1

                a = pkl[labelx][indexx]
                minpos = np.argmin(ma.masked_where(a == 0, a))
                pos_dic = train_set[data_index[cidx][1] + minpos]
                #print(pos_dic[1])
                neg_label = int(labelx)

                while True:
                    neg_label = random.choice(range(1, 770))
                    if neg_label is not int(labelx) and os.path.isdir(
                            os.path.join('D:/datasets/veri-split/train',
                                         strint(neg_label))) is True:
                        break

                negative_label = strint(neg_label)
                neg_cid = pidx[negative_label]
                neg_index = random.choice(range(0, len(pkl[negative_label])))

                neg_dic = train_set[data_index[neg_cid][1] + neg_index]
                trainX[i] = img[i]
                trainX[i + train_loader.batch_size] = pos_dic[0]
                trainX[i + (train_loader.batch_size * 2)] = neg_dic[0]
                trainY[i] = cidx
                trainY[i + train_loader.batch_size] = pos_dic[3]
                trainY[i + (train_loader.batch_size * 2)] = neg_dic[3]

            #print(trainY)
            trainX = trainX.cuda()
            trainY = trainY.cuda()

            score, feat = model(trainX, trainY)
            loss = loss_func(score, feat, trainY)
            loss.backward()
            optimizer.step()
            if 'center' in cfg.MODEL.METRIC_LOSS_TYPE:
                for param in center_criterion.parameters():
                    param.grad.data *= (1. / cfg.SOLVER.CENTER_LOSS_WEIGHT)
                optimizer_center.step()

            acc = (score.max(1)[1] == trainY).float().mean()
            loss_meter.update(loss.item(), img.shape[0])
            acc_meter.update(acc, 1)

            if (n_iter + 1) % log_period == 0:
                logger.info(
                    "Epoch[{}] Iteration[{}/{}] Loss: {:.3f}, Acc: {:.3f}, Base Lr: {:.2e}"
                    .format(epoch, (n_iter + 1), len(train_loader),
                            loss_meter.avg, acc_meter.avg,
                            scheduler.get_lr()[0]))
        end_time = time.time()
        time_per_batch = (end_time - start_time) / (n_iter + 1)
        logger.info(
            "Epoch {} done. Time per batch: {:.3f}[s] Speed: {:.1f}[samples/s]"
            .format(epoch, time_per_batch,
                    train_loader.batch_size / time_per_batch))

        if epoch % checkpoint_period == 0:
            torch.save(
                model.state_dict(),
                os.path.join(cfg.OUTPUT_DIR,
                             cfg.MODEL.NAME + '_{}.pth'.format(epoch)))

        if epoch % eval_period == 0:
            model.eval()
            for n_iter, (img, vid, camid, _, _) in enumerate(val_loader):
                with torch.no_grad():
                    img = img.to(device)
                    feat = model(img)
                    evaluator.update((feat, vid, camid))

            cmc, mAP, _, _, _, _, _ = evaluator.compute()
            logger.info("Validation Results - Epoch: {}".format(epoch))
            logger.info("mAP: {:.1%}".format(mAP))
            for r in [1, 5, 10]:
                logger.info("CMC curve, Rank-{:<3}:{:.1%}".format(
                    r, cmc[r - 1]))