main.py

from __future__ import print_function
import argparse
from math import log10, ceil
import random, shutil, json
from os.path import join, exists, isfile, realpath, dirname
from os import makedirs, remove, chdir, environ

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import DataLoader, SubsetRandomSampler
from torch.utils.data.dataset import Subset
import torchvision.transforms as transforms
from PIL import Image
from datetime import datetime
import torchvision.datasets as datasets
import torchvision.models as models
import h5py
import faiss
from tqdm import tqdm
from tensorboardX import SummaryWriter
import numpy as np
import netvlad
import time, os

import wandb
environ['WANDB_API_KEY'] = ""

from pittsburgh import prefix_data

parser = argparse.ArgumentParser(description='MultiRes-NetVLAD')
parser.add_argument('--mode', type=str, default='train', help='Mode', choices=['train', 'test', 'cluster'])
parser.add_argument('--batchSize', type=int, default=4,
        help='Number of triplets (query, pos, negs). Each triplet consists of 12 images.')
parser.add_argument('--cacheBatchSize', type=int, default=12, help='Batch size for caching and testing')
parser.add_argument('--cacheRefreshRate', type=int, default=1000, 
        help='How often to refresh cache, in number of queries. 0 for off')
parser.add_argument('--nEpochs', type=int, default=30, help='number of epochs to train for')
parser.add_argument('--start_epoch', default=0, type=int, metavar='N', 
        help='manual epoch number (useful on restarts)')
parser.add_argument('--nGPU', type=int, default=1, help='number of GPU to use.')
parser.add_argument('--optim', type=str, default='SGD', help='optimizer to use', choices=['SGD', 'ADAM'])
parser.add_argument('--lr', type=float, default=0.0001, help='Learning Rate.')
parser.add_argument('--lrStep', type=float, default=5, help='Decay LR ever N steps.')
parser.add_argument('--lrGamma', type=float, default=0.5, help='Multiply LR by Gamma for decaying.')
parser.add_argument('--weightDecay', type=float, default=0.001, help='Weight decay for SGD.')
parser.add_argument('--momentum', type=float, default=0.9, help='Momentum for SGD.')
parser.add_argument('--nocuda', action='store_true', default=False, help='Dont use cuda')
parser.add_argument('--threads', type=int, default=8, help='Number of threads for each data loader to use')
parser.add_argument('--seed', type=int, default=123, help='Random seed to use.')
parser.add_argument('--dataPath', type=str, default=prefix_data+'/code_outputs/data/', help='Path for centroid data.')
parser.add_argument('--runsPath', type=str, default=prefix_data+'/code_outputs/runs/', help='Path to save runs to.')
parser.add_argument('--savePath', type=str, default='checkpoints', 
        help='Path to save checkpoints to in logdir. Default=checkpoints/')
parser.add_argument('--cachePath', type=str, default=prefix_data+'/code_outputs/cache/', help='Path to save cache to.')
parser.add_argument('--resume', type=str, default='',\
    help='Path to load checkpoint from, for resuming training or testing.')
parser.add_argument('--appendPcaLayer', action='store_true', default=False, help='Use PCA Whitening for Testing')
parser.add_argument('--ckpt', type=str, default='latest', 
        help='Resume from latest or best checkpoint.', choices=['latest', 'best'])        
parser.add_argument('--evalEvery', type=int, default=1, 
        help='Do a validation set run, and save, every N epochs.')
parser.add_argument('--patience', type=int, default=10, help='Patience for early stopping. 0 is off.')
parser.add_argument('--dataset', type=str, default='pittsburgh', 
        help='Dataset to use', choices=['pittsburgh'])
parser.add_argument('--arch', type=str, default='vgg16-12', 
        help='basenetwork to use, -12 indicates the number of layers to be trained', choices=['vgg16-12'])
parser.add_argument('--trim', type=int, default=-2, help='Layers to be trimmed from end')        
parser.add_argument('--vladv2', action='store_true', help='Use VLAD v2')
parser.add_argument('--pooling', type=str, default='netvlad', help='type of pooling to use',
        choices=['netvlad'])
parser.add_argument('--density_L', type=int, default=10, help='Levels for MultiRes-NetVLAD, e.g., 10=[1,2,3,...,10]')
parser.add_argument('--num_clusters', type=int, default=64, help='Number of NetVlad clusters. Default=64')
parser.add_argument('--margin', type=float, default=0.1, help='Margin for triplet loss. Default=0.1')
parser.add_argument('--nNeg', type=int, default=10, help='Num negs to consider simultaneously for triplet loss. Default=10')
parser.add_argument('--split', type=str, default='val', help='Data split to use for testing. Default is val', 
        choices=['train', 'test', 'val', 'test250k','val250k'])
parser.add_argument('--fromscratch', action='store_true', help='Train from scratch rather than using pretrained models')
parser.add_argument('--imgWidth', type=int, default=640, help='Resize image width')
parser.add_argument('--imgHeight', type=int, default=480, help='Resize image height')
parser.add_argument('--testTrainMem', action='store_true', help='skips caching')

def train(epoch):
    epoch_loss = 0
    startIter = 1 # keep track of batch iter across subsets for logging

    if opt.cacheRefreshRate > 0:
        subsetN = ceil(len(train_set) / opt.cacheRefreshRate)
        #TODO randomise the arange before splitting?
        subsetIdx = np.array_split(np.arange(len(train_set)), subsetN)
    else:
        subsetN = 1
        subsetIdx = [np.arange(len(train_set))]

    nBatches = (len(train_set) + opt.batchSize - 1) // opt.batchSize

    for subIter in range(subsetN):
        print('====> Building Cache %d/%d' %(subIter,subsetN))
        model.eval()
        with h5py.File(train_set.cache, mode='w') as h5: 
            pool_size = encoder_dim
            if opt.pooling.lower() == 'netvlad':
                pool_size *= opt.num_clusters
            h5feat = h5.create_dataset("features", 
                    [len(whole_train_set), pool_size], 
                    dtype=np.float32)
            if not opt.testTrainMem:
                with torch.no_grad():
                    for iteration, (input, indices) in tqdm(enumerate(whole_training_data_loader, 1),total=len(whole_training_data_loader)-1, leave=False):
                        input = input.to(device)
                        image_encoding = model.encoder(input)
                        if iteration<2:
                            print('\n')
                            print('******MultiRes-NetVLAD Caching==> l:%s*****'%str(opt.scaleSpace))
                            print('Scale at:1')
                            print('Input Scaled Shape:',input.shape)
                            print('Output Encoded Shape:',image_encoding.shape)                         
                        image_encoding = image_encoding.view(image_encoding.size(0),image_encoding.size(1),-1)
                        for l in opt.scaleSpace:
                            input_scaled = input[:,:,::l,::l]
                            image_scaled_encoding = model.encoder(input_scaled)
                            if iteration<2: 
                                print('Scale at:',l)
                                print('Input Scaled Shape:',input_scaled.shape)
                                print('Output Encoded Shape:',image_scaled_encoding.shape)                             
                            image_encoding = torch.cat([image_encoding, image_scaled_encoding.view(image_scaled_encoding.size(0),\
                                    image_scaled_encoding.size(1),-1)], dim=2)
                        image_encoding=image_encoding.unsqueeze(3)
                        vlad_encoding = model.pool(image_encoding)
                        if iteration<2:
                            print('input',input.shape)
                            print('image_encoding',image_encoding.shape)
                            print('vlad_encoding',vlad_encoding.shape)     
                            
                        h5feat[indices.detach().numpy(), :] = vlad_encoding.detach().cpu().numpy()
                        del input, image_encoding, vlad_encoding
                        torch.cuda.empty_cache()

        sub_train_set = Subset(dataset=train_set, indices=subsetIdx[subIter])

        training_data_loader = DataLoader(dataset=sub_train_set, num_workers=opt.threads, 
                    batch_size=opt.batchSize, shuffle=True, 
                    collate_fn=dataset.collate_fn, pin_memory=False)

        print('Allocated:', torch.cuda.memory_allocated())
        print('Cached:', torch.cuda.memory_reserved())

        model.train()
        loss = None
        for iteration, (query, positives, negatives, 
                negCounts, indices) in tqdm(enumerate(training_data_loader, startIter),\
                                            total=len(training_data_loader),leave=False):
           # with torch.cuda.amp.autocast():
                # some reshaping to put query, pos, negs in a single (N, 3, H, W) tensor
                # where N = batchSize * (nQuery + nPos + nNeg)
            if query is None: continue # in case we get an empty batch

            B, C, H, W = query.shape
            nNeg = torch.sum(negCounts)
            input = torch.cat([query, positives, negatives])
            del query, positives, negatives

            input = input.to(device)
            image_encoding = model.encoder(input)
            if iteration<2:
                print('\n')
                print('******MultiRes-NetVLAD Training==> l:%s*****'%str(opt.scaleSpace))
                print('Scale at:1')                     
            image_encoding = image_encoding.view(image_encoding.size(0),image_encoding.size(1),-1)
            for l in opt.scaleSpace:
                input_scaled = input[:,:,::l,::l]
                image_scaled_encoding = model.encoder(input_scaled)
                if iteration<2: 
                    print('Scale at:',l)               
                image_encoding = torch.cat([image_encoding, image_scaled_encoding.view(image_scaled_encoding.size(0),\
                        image_scaled_encoding.size(1),-1)], dim=2)
            image_encoding=image_encoding.unsqueeze(3)

            vlad_encoding = model.pool(image_encoding)
        
            vladQ, vladP, vladN = torch.split(vlad_encoding, [B, B, nNeg])
                            
            del input, image_encoding, vlad_encoding
            optimizer.zero_grad()
            
            # calculate loss for each Query, Positive, Negative triplet
            # due to potential difference in number of negatives have to 
            # do it per query, per negative
            loss = 0
            for i, negCount in enumerate(negCounts):
                for n in range(negCount):
                    negIx = (torch.sum(negCounts[:i]) + n).item()
                    loss += criterion(vladQ[i:i+1], vladP[i:i+1], vladN[negIx:negIx+1])

            del vladQ, vladP, vladN
            loss /= nNeg.float().to(device) # normalise by actual number of negatives
            loss.backward()
            optimizer.step()

            batch_loss = loss.item()
            epoch_loss += batch_loss
            if opt.testTrainMem:
                print("Mem test passed!")
                exit(-1)

            if iteration % 50 == 0 or nBatches <= 10:
                print("==> Epoch[{}]({}/{}): Loss: {:.4f}".format(epoch, iteration, 
                    nBatches, batch_loss), flush=True)
                writer.add_scalar('Train/Loss', batch_loss, 
                        ((epoch-1) * nBatches) + iteration)
                writer.add_scalar('Train/nNeg', nNeg, 
                        ((epoch-1) * nBatches) + iteration)
                wandb.log({"loss":batch_loss, "nNeg":nNeg, "epoch":epoch})
                print('Allocated:', torch.cuda.memory_allocated())
                print('Cached:', torch.cuda.memory_reserved())

        startIter += len(training_data_loader)
        del training_data_loader, loss
        optimizer.zero_grad()
        torch.cuda.empty_cache()
        remove(train_set.cache) # delete HDF5 cache

    avg_loss = epoch_loss / nBatches

    print("===> Epoch {} Complete: Avg. Loss: {:.4f}".format(epoch, avg_loss), 
            flush=True)
    wandb.log({"loss_e":avg_loss})
    writer.add_scalar('Train/AvgLoss', avg_loss, epoch)

def test(eval_set, epoch=0, write_tboard=False):
    print("Testing====>")
    print('===>Parameters: ',sum(p.numel() for p in model.parameters()))
    
    # TODO what if features dont fit in memory? 
    test_data_loader = DataLoader(dataset=eval_set, 
                num_workers=opt.threads, batch_size=opt.cacheBatchSize, shuffle=False, 
                pin_memory=False)

    model.eval()
    t_extract_start = time.time()
    with torch.no_grad():
        print('====> Extracting Features')
        pool_size = encoder_dim
        if opt.pooling.lower() == 'netvlad':
            pool_size *= opt.num_clusters
        
        if opt.appendPcaLayer == True and opt.mode.lower() == 'test':
            pool_size = 4096
        dbFeat = np.empty((len(eval_set), pool_size))

        for iteration, (input, indices) in tqdm(enumerate(test_data_loader, 1),total=len(test_data_loader)-1):
            input = input.to(device)
            image_encoding = model.encoder(input)
            if iteration<2:
                print('\n')
                print('******MultiRes-NetVLAD Testing==> l:%s*****'%str(opt.scaleSpace))
                print('Scale at:1')                     
            image_encoding = image_encoding.view(image_encoding.size(0),image_encoding.size(1),-1)            
            for l in opt.scaleSpace:
                input_scaled = input[:,:,::l,::l]
                image_scaled_encoding = model.encoder(input_scaled)
                if iteration<2: 
                    print('Scale at:',l)                  
                image_encoding = torch.cat([image_encoding, image_scaled_encoding.view(image_scaled_encoding.size(0),\
                        image_scaled_encoding.size(1),-1)], dim=2)
            image_encoding=image_encoding.unsqueeze(3)

            vlad_encoding = model.pool(image_encoding) 
            if opt.appendPcaLayer == True and opt.mode.lower() == 'test':
                vlad_encoding = model.WPCA(vlad_encoding.unsqueeze(-1).unsqueeze(-1))
                     
            dbFeat[indices.detach().numpy(), :] = vlad_encoding.detach().cpu().numpy()
            if iteration % 50 == 0 or len(test_data_loader) <= 10:
                print("==> Batch ({}/{})".format(iteration, 
                    len(test_data_loader)), flush=True)

            del input, image_encoding, vlad_encoding
    del test_data_loader
    t_extract_dur = time.time() - t_extract_start

    # extracted for both db and query, now split in own sets
    qFeat = dbFeat[eval_set.dbStruct.numDb:].astype('float32')
    dbFeat = dbFeat[:eval_set.dbStruct.numDb].astype('float32')
    
    print('====> Building faiss index')
    faiss_index = faiss.IndexFlatL2(pool_size)
    faiss_index.add(dbFeat)

    print('====> Calculating recall @ N')
    n_values = [1,5,10,20,50,100]

    t_retrieve_start = time.time()

    _, predictions = faiss_index.search(qFeat, max(n_values)) 
    t_retrieve_dur = time.time() - t_retrieve_start

    # for each query get those within threshold distance
    gt = eval_set.getPositives() 

    correct_at_n = np.zeros(len(n_values))
    #TODO can we do this on the matrix in one go?
    for qIx, pred in enumerate(predictions):
        for i,n in enumerate(n_values):
            # if in top N then also in top NN, where NN > N
            if np.any(np.in1d(pred[:n], gt[qIx])):
                correct_at_n[i:] += 1
                break
    recall_at_n = correct_at_n / eval_set.dbStruct.numQ

    recalls = {} #make dict for output
    for i,n in enumerate(n_values):
        recalls[n] = recall_at_n[i]
        print("====> Recall@{}: {:.4f}".format(n, recall_at_n[i]))
        if write_tboard: writer.add_scalar('Val/Recall@' + str(n), recall_at_n[i], epoch)
        wandb.log({'Val/Recall@' + "{0:03d}".format(n): recall_at_n[i]},commit=False)

    return recalls

def get_clusters(cluster_set):
    nDescriptors = 50000
    nPerImage = 16
    nIm = ceil(nDescriptors/nPerImage)

    sampler = SubsetRandomSampler(np.random.choice(len(cluster_set), nIm, replace=False))
    data_loader = DataLoader(dataset=cluster_set, 
                num_workers=opt.threads, batch_size=opt.cacheBatchSize, shuffle=False, 
                pin_memory=False,
                sampler=sampler)

    if not exists(join(opt.dataPath, 'centroids')):
        makedirs(join(opt.dataPath, 'centroids'))

    initcache = join(opt.dataPath, 'centroids', opt.arch + '_L_'+ str(opt.density_L) + "_" +  opt.pooling + '_' + cluster_set.dataset + '_cluster' + str(opt.num_clusters) + '_trim' + str(opt.trim) + '_desc_cen.hdf5')
    with h5py.File(initcache, mode='w') as h5: 
        with torch.no_grad():
            model.eval()
            print('====> Extracting Descriptors')
            dbFeat = h5.create_dataset("descriptors", 
                        [nDescriptors, encoder_dim], 
                        dtype=np.float32)

            for iteration, (input, indices) in tqdm(enumerate(data_loader, 1),total=len(data_loader)-1):
                input = input.to(device)
                image_encoding = model.encoder(input)
                if iteration<2:
                    print('\n')
                    print('******MultiRes-NetVLAD Clustering: %s*****'%str(opt.scaleSpace))
                    print('Scale at:1')

                image_encoding = image_encoding.view(image_encoding.size(0),image_encoding.size(1),-1)
                
                for l in opt.scaleSpace:
                    input_scaled = input[:,:,::l,::l]
                    image_scaled_encoding = model.encoder(input_scaled)
                    if iteration<2: 
                        print('Scale at:',l)                   
                    image_encoding = torch.cat([image_encoding, image_scaled_encoding.view(image_scaled_encoding.size(0),\
                            image_scaled_encoding.size(1),-1)], dim=2)

                image_descriptors=image_encoding.permute(0, 2, 1)  
#                image_descriptors = model.encoder(input).view(input.size(0), encoder_dim, -1).permute(0, 2, 1)

                batchix = (iteration-1)*opt.cacheBatchSize*nPerImage
                for ix in range(image_descriptors.size(0)):
                    # sample different location for each image in batch
                    sample = np.random.choice(image_descriptors.size(1), min(nPerImage,image_descriptors.size(1)), replace=False)
                    startix = batchix + ix*nPerImage
                    dbFeat[startix:startix+nPerImage, :] = image_descriptors[ix, sample, :].detach().cpu().numpy()

                if iteration % 50 == 0 or len(data_loader) <= 10:
                    print("==> Batch ({}/{})".format(iteration, 
                        ceil(nIm/opt.cacheBatchSize)), flush=True)
                del input, image_descriptors
        
        print('====> Clustering..')
        niter = 100
        kmeans = faiss.Kmeans(encoder_dim, opt.num_clusters, niter=niter, verbose=False)
        kmeans.train(dbFeat[...])

        print('====> Storing centroids', kmeans.centroids.shape)
        h5.create_dataset('centroids', data=kmeans.centroids)
        print('====> Done!')

def save_checkpoint(savePath,state, is_best, filename='checkpoint.pth.tar'):
    model_out_path = join(savePath, filename)
    torch.save(state, model_out_path)
    if is_best:
        shutil.copyfile(model_out_path, join(savePath, 'model_best.pth.tar'))

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class L2Norm(nn.Module):
    def __init__(self, dim=1):
        super().__init__()
        self.dim = dim

    def forward(self, input):
        return F.normalize(input, p=2, dim=self.dim)


if __name__ == "__main__":
    opt = parser.parse_args()

    restore_var = ['lr', 'lrStep', 'lrGamma', 'weightDecay', 'momentum', 
            'runsPath', 'savePath', 'arch', 'num_clusters', 'pooling', 'optim',
            'margin', 'seed', 'patience','trim']
    if opt.resume:
        flag_file = join(opt.resume, 'checkpoints', 'flags.json')
        if exists(flag_file):
            with open(flag_file, 'r') as f:
                stored_flags = {'--'+k : str(v) for k,v in json.load(f).items() if k in restore_var}
                to_del = []
                for flag, val in stored_flags.items():
                    for act in parser._actions:
                        if act.dest == flag[2:]:
                            # store_true / store_false args don't accept arguments, filter these 
                            if type(act.const) == type(True):
                                if val == str(act.default):
                                    to_del.append(flag)
                                else:
                                    stored_flags[flag] = ''
                for flag in to_del: del stored_flags[flag]

                train_flags = [x for x in list(sum(stored_flags.items(), tuple())) if len(x) > 0]
                print('Restored flags:', train_flags)
                opt = parser.parse_args(train_flags, namespace=opt)

    print(opt)

    if opt.pooling == 'netvlad':
        if opt.density_L == 10:
            opt.scaleSpace = [2,3,4,5,6,7,8,9,10]
        elif opt.density_L == 3:
            opt.scaleSpace = [2,4]
            
    wandb_dataStr = opt.dataset.lower()[:4]
    wandbResume = False if opt.resume == '' else True
    wandb.init(project="MultiRes-NetVLAD", entity="ahmedest61",config=opt,resume=wandbResume)
    
    if wandbResume == False:
        wandb_proj_id = wandb.util.generate_id()
        wandb.init(project="MultiRes-NetVLAD",entity="ahmedest61",id=wandb_proj_id,resume="allow")
        opt.resume_id = wandb_proj_id
    else:
        # hard code
        opt.resume_id = "abc"
        wandb.init(project="MultiRes-NetVLAD",entity="ahmedest61",id=opt.resume_id,resume=wandbResume)
        wandb_proj_id = opt.resume_id    
    # update runName
    runName = wandb.run.name

    if opt.resume:
        opt.savePath = join(opt.resume, 'checkpoints')
        writer = SummaryWriter(log_dir=opt.savePath.split('checkpoints')[0])
        # write checkpoints in logdir
        logdir = writer.file_writer.get_logdir()
        dataTime = str(opt.resume).split('/')[-1].split(opt.arch)[0][:-1]
    else:
        dataTime = datetime.now().strftime('%b%d_%H-%M-%S')
    
    expName=dataTime+'_'+opt.arch+'_L_'+ str(opt.density_L) + "_" + opt.pooling + "_" + opt.dataset +"_cluster" + str(opt.num_clusters)+ '_trim' + str(opt.trim)
                                
    if expName != '':
        wandb.run.name = expName + "-" + runName.split("-")[-1]
        wandb.run.save()
    else:
        expName = runName

    opt = wandb.config

    cuda = not opt.nocuda
    if cuda and not torch.cuda.is_available():
        raise Exception("No GPU found, please run with --nocuda")

    device = torch.device("cuda" if cuda else "cpu")

    random.seed(opt.seed)
    np.random.seed(opt.seed)
    torch.manual_seed(opt.seed)
    if cuda:
        torch.cuda.manual_seed(opt.seed)

    print('===> Loading dataset(s)')

    if opt.dataset.lower() == 'pittsburgh':
        import pittsburgh as dataset
        print('pittsburgh')
        if opt.mode.lower() == 'train':
            whole_train_set = dataset.get_whole_training_set()
            whole_training_data_loader = DataLoader(dataset=whole_train_set, 
                    num_workers=opt.threads, batch_size=opt.cacheBatchSize, shuffle=False, 
                    pin_memory=False)

            train_set = dataset.get_training_query_set(opt.margin,opt.nNeg)

            print('====> Training query set:', len(train_set))
            whole_test_set = dataset.get_whole_val_set()
            print('===> Evaluating on val set, query count:', whole_test_set.dbStruct.numQ)
        elif opt.mode.lower() == 'test':

            if opt.split.lower() == 'test':
                whole_test_set = dataset.get_whole_test_set()
                print('===> Evaluating on test30k set')
            elif opt.split.lower() == 'val':
                whole_test_set = dataset.get_whole_val_set()
                print('===> Evaluating on val30k set')
            
            elif opt.split.lower() == 'test250k':
                whole_test_set = dataset.get_250k_test_set()
                print('===> Evaluating on test250k set')
            elif opt.split.lower() == 'val250k':
                whole_test_set = dataset.get_250k_val_set()
                print('===> Evaluating on val250k set')

            elif opt.split.lower() == 'train':
                whole_test_set = dataset.get_whole_training_set()
                print('===> Evaluating on train set')

            else:
                raise ValueError('Unknown dataset split: ' + opt.split)
            print('====> Query count:', whole_test_set.dbStruct.numQ)
        elif opt.mode.lower() == 'cluster':
            whole_train_set = dataset.get_whole_training_set(onlyDB=True)
    else:
        raise Exception('Unknown dataset')
        
    print('===> Building model')

    pretrained = not opt.fromscratch
    if 'vgg16' in opt.arch.lower():
        reqGradLayer = int(opt.arch.split('-')[1])*-1
        encoder = models.vgg16(pretrained=pretrained)
        # capture only feature part and remove last relu and maxpool
        layers = list(encoder.features.children())[:opt.trim]
        encoder_dim = layers[-1].out_channels 
        if pretrained:
            # if using pretrained then 
            for l in layers[:reqGradLayer]: 
                for p in l.parameters():
                    p.requires_grad = False

    print("\nEncoder Dim: ",encoder_dim)

    if opt.mode.lower() == 'cluster' and not opt.vladv2:
        layers.append(L2Norm())

    encoder = nn.Sequential(*layers)
    model = nn.Module() 
    model.add_module('encoder', encoder)

    if opt.mode.lower() != 'cluster':
        if opt.pooling.lower() == 'netvlad':
            net_vlad = netvlad.NetVLAD(num_clusters=opt.num_clusters, dim=encoder_dim, vladv2=opt.vladv2)
            if not opt.resume: 
                if opt.mode.lower() == 'train':
                    initcache = join(opt.dataPath, 'centroids', opt.arch + '_L_'+ str(opt.density_L) + "_" +  opt.pooling + '_' + train_set.dataset + '_cluster' + str(opt.num_clusters) + '_trim' + str(opt.trim) +'_desc_cen.hdf5')
                else:
                    initcache = join(opt.dataPath, 'centroids', opt.arch + '_L_'+ str(opt.density_L) + "_" +  opt.pooling + '_' + whole_test_set.dataset + '_cluster' + str(opt.num_clusters) + '_trim' + str(opt.trim) +'_desc_cen.hdf5')

                if not exists(initcache):
                    raise FileNotFoundError('Could not find clusters, please run with --mode=cluster before proceeding')
                    print(initcache)
                else:
                    print("InitCache Loaded==>")
                    print(initcache)
                with h5py.File(initcache, mode='r') as h5: 
                    clsts = h5.get("centroids")[...]
                    traindescs = h5.get("descriptors")[...]
                    net_vlad.init_params(clsts, traindescs) 
                    del clsts, traindescs

            model.add_module('pool', net_vlad)

        else:
            raise ValueError('Unknown pooling type: ' + opt.pooling)

        if opt.appendPcaLayer == True: # should be True for testing only 
            if opt.mode.lower() == 'test':
                num_pcs = 4096
                netvlad_output_dim = encoder_dim
                if opt.pooling.lower() == 'netvlad':
                    netvlad_output_dim *= opt.num_clusters
            
                pca_conv = nn.Conv2d(netvlad_output_dim, num_pcs, kernel_size=(1, 1), stride=1, padding=0)
                model.add_module('WPCA', nn.Sequential(*[pca_conv, Flatten(), L2Norm(dim=-1)]))
                
                if opt.ckpt.lower() == 'latest':
                    checkpoint_tar = 'checkpoint_WPCA_4096.pth.tar'
                print("==>WPCA Added")
            elif opt.mode.lower() == 'train' or opt.mode.lower() == 'cluster':
                raise ValueError('pca_layer can be used only during Testing')        

    print(model)

    isParallel = False
    if opt.nGPU > 1 and torch.cuda.device_count() > 1:
        model.encoder = nn.DataParallel(model.encoder)
        if opt.mode.lower() != 'cluster':
            model.pool = nn.DataParallel(model.pool)
        isParallel = True

    if not opt.resume:
        model = model.to(device)
    
    if opt.mode.lower() == 'train':
        if opt.optim.upper() == 'ADAM':
            optimizer = optim.Adam(filter(lambda p: p.requires_grad, 
                model.parameters()), lr=opt.lr)#, betas=(0,0.9))
        elif opt.optim.upper() == 'SGD':
            optimizer = optim.SGD(filter(lambda p: p.requires_grad, 
                model.parameters()), lr=opt.lr,
                momentum=opt.momentum,
                weight_decay=opt.weightDecay)

            scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=opt.lrStep, gamma=opt.lrGamma)
        else:
            raise ValueError('Unknown optimizer: ' + opt.optim)

        # original paper/code doesn't sqrt() the distances, we do, so sqrt() the margin, I think :D
        criterion = nn.TripletMarginLoss(margin=opt.margin**0.5, 
                p=2, reduction='sum').to(device)

    if opt.resume:
        resume_ckpt = join(opt.resume, 'checkpoints', checkpoint_tar)
        if opt.appendPcaLayer == False:
            if opt.ckpt.lower() == 'latest':
                checkpoint_tar = 'checkpoint.pth.tar'
            elif opt.ckpt.lower() == 'best':
                checkpoint_tar = 'model_best.pth.tar' 
        if isfile(resume_ckpt):
            print("=> loading checkpoint '{}'".format(resume_ckpt))
            checkpoint = torch.load(resume_ckpt, map_location=lambda storage, loc: storage)
            if opt.appendPcaLayer == False:
                start_epoch = checkpoint['epoch']
                best_metric = checkpoint['best_score']
            model.load_state_dict(checkpoint['state_dict'])
            if opt.appendPcaLayer == True:
                assert (checkpoint['num_pcs']==4096)
            model = model.to(device)
            if opt.mode == 'train':
                optimizer.load_state_dict(checkpoint['optimizer'])
            if opt.appendPcaLayer == False:
                print("=> loaded checkpoint '{}' (epoch {})"
                    .format(resume_ckpt, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(resume_ckpt))
    else:
        start_epoch = opt.start_epoch
    if opt.mode.lower() == 'test':
        print('===> Running evaluation step')
        epoch = 1
        recalls = test(whole_test_set, epoch, write_tboard=False)
    elif opt.mode.lower() == 'cluster':
        print('===> Calculating descriptors and clusters')
        get_clusters(whole_train_set)
    elif opt.mode.lower() == 'train':
        print('===> Training model')
        writer = SummaryWriter(log_dir=join(opt.runsPath,expName))
        timeStampStr = datetime.now().strftime('%b%d_%H-%M-%S-%f')
        train_set.cache = join(opt.cachePath, timeStampStr +'_'+train_set.whichSet + '_L_'+ str(opt.density_L) + "_" +  opt.pooling+ "_cluster" + str(opt.num_clusters)+ '_trim' + str(opt.trim)+'_feat_cache.hdf5')        
        
        # write checkpoints in logdir
        logdir = writer.file_writer.get_logdir()
        savePath = join(logdir, opt.savePath)
        if not opt.resume:
            makedirs(savePath)

        with open(join(savePath, 'flags.json'), 'w') as f:
            f.write(json.dumps({k:v for k,v in opt.items()}))

        print('===> Saving state to:', logdir)

        not_improved = 0
        best_score = 0
        for epoch in range(start_epoch+1, opt.nEpochs + 1):
            train(epoch)
            if opt.optim.upper() == 'SGD':
                scheduler.step(epoch)
            if (epoch % opt.evalEvery) == 0:
                recalls = test(whole_test_set, epoch, write_tboard=True)
                is_best = recalls[5] > best_score 
                if is_best:
                    not_improved = 0
                    best_score = recalls[5]
                else: 
                    not_improved += 1

                save_checkpoint(savePath,{
                        'epoch': epoch,
                        'state_dict': model.state_dict(),
                        'recalls': recalls,
                        'best_score': best_score,
                        'optimizer' : optimizer.state_dict(),
                        'parallel' : isParallel,
                }, is_best)

                if opt.patience > 0 and not_improved > (opt.patience / opt.evalEvery):
                    print('Performance did not improve for', opt.patience, 'epochs. Stopping.')
                    break

        print("=> Best Recall@5: {:.4f}".format(best_score), flush=True)
        writer.close()