def loadModelToCuda(eConfig):
kargs = {
'model': eConfig['arch'],
'pooling': eConfig['pooling'],
'whitening': False
} #eConfig['whitening'] }
if (eConfig['pretrained'] == 'imageNet'):
model = init_network(**kargs, pretrained=True)
else:
model = init_network(**kargs, pretrained=False)
return model.cuda()
def constructfeature(self, hash_size, input_dim, num_hashtables):
multiscale = '[1]'
print(">> Loading network:\n>>>> '{}'".format(self.network))
# state = load_url(PRETRAINED[args.network], model_dir=os.path.join(get_data_root(), 'networks'))
state = torch.load(self.network)
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get('local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
# network initialization
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = list(eval(multiscale))
print(">>>> Evaluating scales: {}".format(ms))
# moving network to gpu and eval mode
if torch.cuda.is_available():
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(
mean=net.meta['mean'],
std=net.meta['std']
)
transform = transforms.Compose([
transforms.ToTensor(),
normalize
])
# extract database and query vectors
print('>> database images...')
images = ImageProcess(self.img_dir).process()
vecs, img_paths = extract_vectors(net, images, 1024, transform, ms=ms)
feature_dict = dict(zip(img_paths, list(vecs.detach().cpu().numpy().T)))
# index
lsh = LSHash(hash_size=int(hash_size), input_dim=int(input_dim), num_hashtables=int(num_hashtables))
for img_path, vec in feature_dict.items():
lsh.index(vec.flatten(), extra_data=img_path)
# ## 保存索引模型
# with open(self.feature_path, "wb") as f:
# pickle.dump(feature_dict, f)
# with open(self.index_path, "wb") as f:
# pickle.dump(lsh, f)
print("extract feature is done")
return feature_dict, lsh
def load_network(model_dir: str, device: torch.device) -> ImageRetrievalNet:
state = load_url(PRETRAINED['rSfM120k-tl-resnet101-gem-w'],
model_dir=model_dir)
net = init_network({
'architecture': state['meta']['architecture'],
'pooling': state['meta']['pooling'],
'whitening': state['meta'].get('whitening', False)
})
net.load_state_dict(state['state_dict'])
net.to(device)
net.eval()
log.info(f"Loaded network: {net.meta_repr()}")
return net
def test_feature(self):
multiscale = '[1]'
print(">> Loading network:\n>>>> '{}'".format(self.network))
# state = load_url(PRETRAINED[args.network], model_dir=os.path.join(get_data_root(), 'networks'))
state = torch.load(self.network)
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get('local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
# network initialization
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = list(eval(multiscale))
print(">>>> Evaluating scales: {}".format(ms))
# moving network to gpu and eval mode
if torch.cuda.is_available():
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(
mean=net.meta['mean'],
std=net.meta['std']
)
transform = transforms.Compose([
transforms.ToTensor(),
normalize
])
# extract database and query vectors
print('>> database images...')
images = ImageProcess(self.img_dir).process()
vecs, img_paths = extract_vectors(net, images, 1024, transform, ms=ms)
feature_dict = dict(zip(img_paths, list(vecs.detach().cpu().numpy().T)))
return feature_dict
def load_offtheshelf(network_name):
"""Load off the shelf network."""
offtheshelf = network_name.split('-')
if len(offtheshelf)==3:
if offtheshelf[2]=='whiten':
offtheshelf_whiten = True
else:
raise(RuntimeError("Incorrect format of the off-the-shelf network. Examples: resnet101-gem | resnet101-gem-whiten"))
else:
offtheshelf_whiten = False
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(network_name))
net = init_network(model=offtheshelf[0], pooling=offtheshelf[1], whitening=offtheshelf_whiten)
print(">>>> loaded network: ")
print(net.meta_repr())
return net
def load_network(network_path):
"""Load network from a path."""
print(">> Loading network:\n>>>> '{}'".format(network_path))
if network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[network_path], model_dir=os.path.join(get_data_root(), 'networks'))
else:
state = torch.load(network_path)
net = init_network(model=state['meta']['architecture'], pooling=state['meta']['pooling'], whitening=state['meta']['whitening'],
mean=state['meta']['mean'], std=state['meta']['std'], pretrained=False)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
return net
def build_retrievalnet_from_options(opt, is_cuda=False):
logger = logging.getLogger("detector-retrieval.build_retrievalnet")
logger.info("Building the retrieval model...")
logger.info("Loading weights from {}".format(opt.retrieval_network_path))
state = torch.load(opt.retrieval_network_path)
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params["architecture"] = state["meta"]["architecture"]
net_params["pooling"] = state["meta"]["pooling"]
net_params["local_whitening"] = state["meta"].get("local_whitening", False)
net_params["regional"] = state["meta"].get("regional", False)
net_params["whitening"] = state["meta"].get("whitening", False)
net_params["mean"] = state["meta"]["mean"]
net_params["std"] = state["meta"]["std"]
net_params["pretrained"] = False
# load network
net = init_network(net_params)
net.load_state_dict(state["state_dict"])
# if whitening is precomputed
if "Lw" in state["meta"]:
net.meta["Lw"] = state["meta"]["Lw"]
logger.info("Loaded network: ")
if "epoch" in state:
logger.info("Model after {} epochs".format(state["epoch"]))
logger.info(net.meta_repr())
if is_cuda:
net.cuda()
else:
net.cpu()
net.eval()
return net
def main():
args = parser.parse_args()
# check if there are unknown datasets
for dataset in args.datasets.split(','):
if dataset not in datasets_names:
raise ValueError(
'Unsupported or unknown dataset: {}!'.format(dataset))
# check if test dataset are downloaded
# and download if they are not
download_train(get_data_root())
download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network
# pretrained networks (downloaded automatically)
print(">> Loading network:\n>>>> '{}'".format(args.network))
state = load_url(PRETRAINED[args.network],
model_dir=os.path.join(get_data_root(), 'networks'))
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get('local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
# network initialization
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = list(eval(args.multiscale))
print(">>>> Evaluating scales: {}".format(ms))
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(mean=net.meta['mean'],
std=net.meta['std'])
transform = transforms.Compose([transforms.ToTensor(), normalize])
# evaluate on test datasets
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
print('>> {}: Extracting...'.format(dataset))
# prepare config structure for the test dataset
cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])]
qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])]
try:
bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
except:
bbxs = None # for holidaysmanrot and copydays
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net, images, args.image_size, transform, ms=ms)
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net,
qimages,
args.image_size,
transform,
bbxs=bbxs,
ms=ms)
print('>> {}: Evaluating...'.format(dataset))
# convert to numpy
vecs = vecs.numpy()
qvecs = qvecs.numpy()
# search, rank, and print
scores = np.dot(vecs.T, qvecs)
ranks = np.argsort(-scores, axis=0)
top_k = 100
ranks_fnames_qs = []
for q_id in range(len(cfg["qimlist"])):
ranks_q = list(ranks[:top_k, q_id])
ranks_fname_per_q = []
for img_id in ranks_q:
ranks_fname_per_q.append(cfg["imlist"][img_id])
ranks_fnames_qs.append(ranks_fname_per_q)
compute_map_and_print(dataset, ranks, cfg['gnd'])
compute_map_and_print_top_k(dataset, ranks_fnames_qs, cfg['gnd'],
cfg["imlist"])
sys.exit()
with open(dataset + "_gl18_tl_resnet101_gem_w_m.pkl", "wb") as f:
data = {"ranks": ranks, "db_images": images, "q_images": qimages}
pickle.dump(data, f)
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
def main():
args = parser.parse_args()
# check if test dataset are downloaded
# and download if they are not
download_train(get_data_root())
download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
state = torch.load(args.network_path)
net = init_network(model=state['meta']['architecture'],
pooling=state['meta']['pooling'],
whitening=state['meta']['whitening'],
mean=state['meta']['mean'],
std=state['meta']['std'],
pretrained=False)
net.load_state_dict(state['state_dict'])
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
offtheshelf = args.network_offtheshelf.split('-')
if len(offtheshelf) == 3:
if offtheshelf[2] == 'whiten':
offtheshelf_whiten = True
else:
raise (RuntimeError(
"Incorrect format of the off-the-shelf network. Examples: resnet101-gem | resnet101-gem-whiten"
))
else:
offtheshelf_whiten = False
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(model=offtheshelf[0],
pooling=offtheshelf[1],
whitening=offtheshelf_whiten)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = [1]
msp = 1
if args.multiscale:
ms = [1, 1. / math.sqrt(2), 1. / 2]
if net.meta['pooling'] == 'gem' and net.whiten is None:
msp = net.pool.p.data.tolist()[0]
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(mean=net.meta['mean'],
std=net.meta['std'])
transform = transforms.Compose([transforms.ToTensor(), normalize])
# compute whitening
if args.whitening is not None:
start = time.time()
print('>> {}: Learning whitening...'.format(args.whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train', args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(args.whitening))
with open(db_fn, 'rb') as f:
db = pickle.load(f)
images = [
cid2filename(db['cids'][i], ims_root)
for i in range(len(db['cids']))
]
# extract whitening vectors
print('>> {}: Extracting...'.format(args.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
print('>> {}: elapsed time: {}'.format(args.whitening,
htime(time.time() - start)))
else:
Lw = None
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
print('>> {}: Extracting...'.format(dataset))
if dataset == 'reco':
images, qimages = landmark_recognition_dataset()
bbxs = [None for x in qimages]
elif dataset == 'retr':
images, _ = landmark_retrieval_dataset()
qimages = []
bbxs = [None for x in qimages]
else:
# prepare config structure for the test dataset
cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])]
qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])]
bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
with open('%s_fnames.pkl' % dataset, 'wb') as f:
pickle.dump([images, qimages], f)
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
vecs = vecs.numpy()
print('>> saving')
np.save('{}_vecs.npy'.format(dataset), vecs)
if len(qimages) > 0:
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net,
qimages,
args.image_size,
transform,
bbxs=bbxs,
ms=ms,
msp=msp)
qvecs = qvecs.numpy()
np.save('{}_qvecs.npy'.format(dataset), qvecs)
if Lw is not None:
# whiten the vectors
vecs_lw = whitenapply(vecs, Lw['m'], Lw['P'])
qvecs_lw = whitenapply(qvecs, Lw['m'], Lw['P'])
# TODO
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
def main():
global args, min_loss
args = parser.parse_args()
# manually check if there are unknown test datasets
for dataset in args.test_datasets.split(','):
if dataset not in test_datasets_names:
raise ValueError('Unsupported or unknown test dataset: {}!'.format(dataset))
# check if test dataset are downloaded
# and download if they are not
download_train(get_data_root())
download_test(get_data_root())
# create export dir if it doesnt exist
directory = "{}".format(args.training_dataset)
directory += "_{}".format(args.arch)
directory += "_{}".format(args.pool)
if args.local_whitening:
directory += "_lwhiten"
if args.regional:
directory += "_r"
if args.whitening:
directory += "_whiten"
if not args.pretrained:
directory += "_notpretrained"
directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin)
directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr, args.weight_decay)
directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num, args.query_size, args.pool_size)
directory += "_bsize{}_uevery{}_imsize{}".format(args.batch_size, args.update_every, args.image_size)
args.directory = os.path.join(args.directory, directory)
print(">> Creating directory if it does not exist:\n>> '{}'".format(args.directory))
if not os.path.exists(args.directory):
os.makedirs(args.directory)
# set cuda visible device
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# set random seeds
# TODO: maybe pass as argument in future implementation?
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
np.random.seed(0)
# initialize model
if args.pretrained:
print(">> Using pre-trained model '{}'".format(args.arch))
else:
print(">> Using model from scratch (random weights) '{}'".format(args.arch))
model_params = {}
model_params['architecture'] = args.arch
model_params['pooling'] = args.pool
model_params['local_whitening'] = args.local_whitening
model_params['regional'] = args.regional
model_params['whitening'] = args.whitening
# model_params['mean'] = ... # will use default
# model_params['std'] = ... # will use default
model_params['pretrained'] = args.pretrained
model = init_network(model_params)
# move network to gpu
model.cuda()
if model_params['pooling'] == 'netvlad':
normalize = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
transform = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=(364, 364)),
torchvision.transforms.ToTensor(),
normalize,
])
image_folder = folder.ImageFolder(root="/mnt/m2/dataset", transform=transform)
train_loader = torch.utils.data.DataLoader(
image_folder,
batch_size=64,
num_workers=8,
shuffle=True
)
n_batches = 10
descs_list = []
i = 0
with torch.no_grad():
for x, _ in train_loader:
model.eval()
desc = model.compute_features(x.cuda())
max_pooled_feat_3 = torch.nn.functional.max_pool2d(desc, kernel_size=3, stride=1)
max_pooled_feat_2 = torch.nn.functional.max_pool2d(desc, kernel_size=2, stride=1)
reshaped_pool_3 = make_locals(max_pooled_feat_3)
reshaped_pool_2 = make_locals(max_pooled_feat_2)
desc = torch.cat([reshaped_pool_2, reshaped_pool_3], dim=1)
# N, dim, h, w = desc.shape
# desc = desc.view(N, dim, h*w).permute(0, 2, 1).reshape(N, -1, 512)
desc = desc.cpu().numpy().astype('float32')
descs_list.append(desc)
print(">> Extracted batch {}/{} - NetVLAD initialization -".format(i+1, n_batches))
i+=1
if i == n_batches:
break
descs_list = np.array(descs_list).reshape(-1, 512)
print(descs_list.shape)
print(">> Sampling local features ")
# locals = np.vstack((m[np.random.randint(len(m), size=150)] for m in descs_list)).astype('float32')
locals = descs_list[np.random.randint(len(descs_list), size=len(descs_list)//3)]
np.random.shuffle(locals)
print(">> Locals extracted shape : {}".format(locals.shape))
n_clust = 64
locals = preprocessing.normalize(locals, axis=1)
print(">> Fitting centroids with K-Means")
kmeans = MiniBatchKMeans(n_clusters=n_clust).fit(locals)
centroids = kmeans.cluster_centers_
print(">> Centroids shape: ", centroids.shape)
model.pool.init_params(centroids.T)
print(">> NetVLAD initialized")
# define loss function (criterion) and optimizer
if args.loss == 'contrastive':
criterion = ContrastiveLoss(margin=args.loss_margin).cuda()
elif args.loss == 'triplet':
criterion = TripletLoss(margin=args.loss_margin).cuda()
else:
raise(RuntimeError("Loss {} not available!".format(args.loss)))
# parameters split into features, pool, whitening
# IMPORTANT: no weight decay for pooling parameter p in GeM or regional-GeM
parameters = []
# add feature parameters
parameters.append({'params': model.features.parameters()})
# add local whitening if exists
if model.lwhiten is not None:
parameters.append({'params': model.lwhiten.parameters()})
# add pooling parameters (or regional whitening which is part of the pooling layer!)
if not args.regional:
# global, only pooling parameter p weight decay should be 0
if args.pool == 'gem':
parameters.append({'params': model.pool.parameters(), 'lr': args.lr*10, 'weight_decay': 0})
elif args.pool == 'gemmp':
parameters.append({'params': model.pool.parameters(), 'lr': args.lr*100, 'weight_decay': 0})
else:
# regional, pooling parameter p weight decay should be 0,
# and we want to add regional whitening if it is there
if args.pool == 'gem':
parameters.append({'params': model.pool.rpool.parameters(), 'lr': args.lr*10, 'weight_decay': 0})
elif args.pool == 'gemmp':
parameters.append({'params': model.pool.rpool.parameters(), 'lr': args.lr*100, 'weight_decay': 0})
if model.pool.whiten is not None:
parameters.append({'params': model.pool.whiten.parameters()})
# add final whitening if exists
if model.whiten is not None:
parameters.append({'params': model.whiten.parameters()})
# define optimizer
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(parameters, args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(parameters, args.lr, weight_decay=args.weight_decay)
# define learning rate decay schedule
# TODO: maybe pass as argument in future implementation?
exp_decay = math.exp(-0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay)
# optionally resume from a checkpoint
start_epoch = 0
if args.resume:
args.resume = os.path.join(args.directory, args.resume)
if os.path.isfile(args.resume):
# load checkpoint weights and update model and optimizer
print(">> Loading checkpoint:\n>> '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
min_loss = checkpoint['min_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
# important not to forget scheduler updating
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch']-1)
else:
print(">> No checkpoint found at '{}'".format(args.resume))
# Data loading code
normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std'])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
train_dataset = TuplesDataset(
name=args.training_dataset,
mode='train',
imsize=args.image_size,
nnum=args.neg_num,
qsize=args.query_size,
poolsize=args.pool_size,
transform=transform
)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True, sampler=None,
drop_last=True, collate_fn=collate_tuples
)
if args.val:
val_dataset = TuplesDataset(
name=args.training_dataset,
mode='val',
imsize=args.image_size,
nnum=args.neg_num,
qsize=float('Inf'),
poolsize=float('Inf'),
transform=transform
)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True,
drop_last=True, collate_fn=collate_tuples
)
# evaluate the network before starting
# this might not be necessary?
test(args.test_datasets, model)
for epoch in range(start_epoch, args.epochs):
# set manual seeds per epoch
np.random.seed(epoch)
torch.manual_seed(epoch)
torch.cuda.manual_seed_all(epoch)
# adjust learning rate for each epoch
scheduler.step()
# # debug printing to check if everything ok
# lr_feat = optimizer.param_groups[0]['lr']
# lr_pool = optimizer.param_groups[1]['lr']
# print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool))
# train for one epoch on train set
loss = train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
if args.val:
with torch.no_grad():
loss = validate(val_loader, model, criterion, epoch)
# evaluate on test datasets every test_freq epochs
if (epoch + 1) % args.test_freq == 0:
with torch.no_grad():
test(args.test_datasets, model)
# remember best loss and save checkpoint
is_best = loss < min_loss
min_loss = min(loss, min_loss)
save_checkpoint({
'epoch': epoch + 1,
'meta': model.meta,
'state_dict': model.state_dict(),
'min_loss': min_loss,
'optimizer' : optimizer.state_dict(),
}, is_best, args.directory)
def main():
global args, max_meter
args = parser.parse_args()
# create export dir if it doesnt exist
directory = "{}".format(args.training_dataset)
directory += "_{}".format(args.arch)
directory += "_{}".format(args.pool)
if args.local_whitening:
directory += "_lwhiten"
if args.regional:
directory += "_r"
if args.whitening:
directory += "_whiten"
if not args.pretrained:
directory += "_notpretrained"
directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin)
directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr,
args.weight_decay)
directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num,
args.query_size,
args.pool_size)
directory += "_bsize{}_imsize{}".format(args.batch_size, args.image_size)
args.directory = os.path.join(args.directory, directory)
print(">> Creating directory if it does not exist:\n>> '{}'".format(
args.directory))
if not os.path.exists(args.directory):
os.makedirs(args.directory)
# set random seeds
# TODO: maybe pass as argument in future implementation?
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
np.random.seed(0)
# initialize model
if args.pretrained:
print(">> Using pre-trained model '{}'".format(args.arch))
else:
print(">> Using model from scratch (random weights) '{}'".format(
args.arch))
model_params = {}
model_params['architecture'] = args.arch
model_params['pooling'] = args.pool
model_params['local_whitening'] = args.local_whitening
model_params['regional'] = args.regional
model_params['whitening'] = args.whitening
# model_params['mean'] = ... # will use default
# model_params['std'] = ... # will use default
model_params['pretrained'] = args.pretrained
model = init_network(model_params)
# move network to gpu
model.cuda()
# define loss function (criterion) and optimizer
if args.loss == 'contrastive':
criterion = ContrastiveLoss(margin=args.loss_margin).cuda()
else:
raise (RuntimeError("Loss {} not available!".format(args.loss)))
# parameters split into features, pool, whitening
# IMPORTANT: no weight decay for pooling parameter p in GeM or regional-GeM
parameters = []
# add feature parameters
parameters.append({'params': model.features.parameters()})
# add local whitening if exists
if model.lwhiten is not None:
parameters.append({'params': model.lwhiten.parameters()})
# add pooling parameters (or regional whitening which is part of the pooling layer!)
if not args.regional:
# global, only pooling parameter p weight decay should be 0
parameters.append({
'params': model.pool.parameters(),
'lr': args.lr * 10,
'weight_decay': 0
})
else:
# regional, pooling parameter p weight decay should be 0,
# and we want to add regional whitening if it is there
parameters.append({
'params': model.pool.rpool.parameters(),
'lr': args.lr * 10,
'weight_decay': 0
})
if model.pool.whiten is not None:
parameters.append({'params': model.pool.whiten.parameters()})
# add final whitening if exists
if model.whiten is not None:
parameters.append({'params': model.whiten.parameters()})
# define optimizer
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(parameters,
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(parameters,
args.lr,
weight_decay=args.weight_decay)
# define learning rate decay schedule
# TODO: maybe pass as argument in future implementation?
exp_decay = math.exp(-0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer,
gamma=exp_decay)
# optionally resume from a checkpoint
start_epoch = 0
if args.resume:
args.resume = os.path.join(args.directory, args.resume)
if os.path.isfile(args.resume):
# load checkpoint weights and update model and optimizer
print(">> Loading checkpoint:\n>> '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
max_meter = checkpoint['max_meter']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# important not to forget scheduler updating
scheduler = torch.optim.lr_scheduler.ExponentialLR(
optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch'] - 1)
else:
print(">> No checkpoint found at '{}'".format(args.resume))
# Data loading code
normalize = transforms.Normalize(mean=model.meta['mean'],
std=model.meta['std'])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
train_dataset = TuplesDataset(name=args.training_dataset,
mode='train',
imsize=args.image_size,
nnum=args.neg_num,
qsize=args.query_size,
poolsize=args.pool_size,
transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
sampler=None,
drop_last=True,
collate_fn=collate_tuples)
if args.val:
val_dataset = TuplesDataset(name=args.training_dataset,
mode='val',
imsize=args.image_size,
nnum=args.neg_num,
qsize=float('Inf'),
poolsize=float('Inf'),
transform=transform)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
drop_last=True,
collate_fn=collate_tuples)
for epoch in range(start_epoch, args.epochs):
# set manual seeds per epoch
np.random.seed(epoch)
torch.manual_seed(epoch)
torch.cuda.manual_seed_all(epoch)
# evaluate on validation set
if args.val:
with torch.no_grad():
loss = validate(val_loader, model, criterion, epoch)
# evaluate on test datasets every test_freq epochs
if epoch % args.test_freq == 0:
with torch.no_grad():
cur_meter = test(args.test_datasets,
model,
image_size=args.image_size)
# remember best loss and save checkpoint
is_best = cur_meter > max_meter
max_meter = max(cur_meter, max_meter)
save_checkpoint(
{
'epoch': epoch,
'meta': model.meta,
'state_dict': model.state_dict(),
'max_meter': max_meter,
'optimizer': optimizer.state_dict(),
},
is_best,
args.directory,
save_regular=args.save_freq and epoch % args.save_freq == 0
or epoch == 0)
# train for one epoch on train set
loss = train(train_loader, model, criterion, optimizer, epoch)
# adjust learning rate for each epoch
scheduler.step()
def main():
args = parser.parse_args()
# check if there are unknown datasets
for dataset in args.datasets.split(','):
if dataset not in datasets_names:
raise ValueError('Unsupported or unknown dataset: {}!'.format(dataset))
# check if test dataset are downloaded
# and download if they are not
# download_train(get_data_root())
# download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path], model_dir=os.path.join(get_data_root(), 'networks'))
else:
# fine-tuned network from path
state = torch.load(args.network_path)
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get('local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
net_params['multi_layer_cat'] = state['meta']['multi_layer_cat']
# load network
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
# parse off-the-shelf parameters
offtheshelf = args.network_offtheshelf.split('-')
net_params = {}
net_params['architecture'] = offtheshelf[0]
net_params['pooling'] = offtheshelf[1]
net_params['local_whitening'] = 'lwhiten' in offtheshelf[2:]
net_params['regional'] = 'reg' in offtheshelf[2:]
net_params['whitening'] = 'whiten' in offtheshelf[2:]
net_params['pretrained'] = True
# load off-the-shelf network
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(args.network_offtheshelf))
net = init_network(net_params)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
print(">> image size: {}".format(args.image_size))
ms = list(eval(args.multiscale))
if len(ms)>1 and net.meta['pooling'] == 'gem' and not net.meta['regional'] and not net.meta['whitening']:
msp = net.pool.p.item()
print(">> Set-up multiscale:")
print(">>>> ms: {}".format(ms))
print(">>>> msp: {}".format(msp))
else:
msp = 1
print(">> Set-up multiscale:")
print(">>>> ms: {}".format(ms))
print(">>>> msp: {}".format(msp))
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(
mean=net.meta['mean'],
std=net.meta['std']
)
transform = transforms.Compose([
transforms.ToTensor(),
normalize
])
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(args.whitening))
if len(ms)>1:
Lw = net.meta['Lw'][args.whitening]['ms']
else:
Lw = net.meta['Lw'][args.whitening]['ss']
else:
# if we evaluate networks from path we should save/load whitening
# not to compute it every time
if args.network_path is not None:
whiten_fn = args.network_path + '_{}_whiten'.format(args.whitening)
if len(ms) > 1:
whiten_fn += '_ms'
whiten_fn += '.pth'
else:
whiten_fn = None
if whiten_fn is not None and os.path.isfile(whiten_fn):
print('>> {}: Whitening is precomputed, loading it...'.format(args.whitening))
Lw = torch.load(whiten_fn)
else:
print('>> {}: Learning whitening...'.format(args.whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train', args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(args.whitening))
with open(db_fn, 'rb') as f:
db = pickle.load(f)
images = [cid2filename(db['cids'][i], ims_root) for i in range(len(db['cids']))]
# extract whitening vectors
print('>> {}: Extracting...'.format(args.whitening))
wvecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
# saving whitening if whiten_fn exists
if whiten_fn is not None:
print('>> {}: Saving to {}...'.format(args.whitening, whiten_fn))
torch.save(Lw, whiten_fn)
print('>> {}: elapsed time: {}'.format(args.whitening, htime(time.time()-start)))
else:
Lw = None
# evaluate on test datasets
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
print('>> {}: Extracting...'.format(dataset))
# prepare config structure for the test dataset
cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
images = [cfg['im_fname'](cfg,i) for i in range(cfg['n'])]
qimages = [cfg['qim_fname'](cfg,i) for i in range(cfg['nq'])]
# bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
print('>> not use bbxs...')
bbxs = None
# key_url_list = ParseData(os.path.join(get_data_root(), 'index.csv'))
# index_image_path = os.path.join(get_data_root(), 'resize_index_image')
# images = [os.path.join(index_image_path, key_url_list[i][0]) for i in range(len(key_url_list))]
# key_url_list = ParseData(os.path.join(get_data_root(), 'test.csv'))
# test_image_path = os.path.join(get_data_root(), 'resize_test_image')
# qimages = [os.path.join(test_image_path, key_url_list[i][0]) for i in range(len(key_url_list))]
# # bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
# csvfile = open(os.path.join(get_data_root(), 'index_clear.csv'), 'r')
# csvreader = csv.reader(csvfile)
# images = [line[:1][0] for line in csvreader]
#
# csvfile = open(os.path.join(get_data_root(), 'test_clear.csv'), 'r')
# csvreader = csv.reader(csvfile)
# qimages = [line[:1][0] for line in csvreader]
# bbxs = None
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp)
# vecs = torch.randn(2048, 5063)
# vecs = torch.randn(2048, 4993)
# hxq modified
# bbxs = None
# print('>> set no bbxs...')
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net, qimages, args.image_size, transform, bbxs=bbxs, ms=ms, msp=msp)
print('>> {}: Evaluating...'.format(dataset))
# convert to numpy
vecs = vecs.numpy()
qvecs = qvecs.numpy()
# search, rank, and print
scores = np.dot(vecs.T, qvecs)
# hxq modified, test add features map for retrieval
# vecs = [vecs[i].numpy() for i in range(len(vecs))]
# qvecs_temp = np.zeros((qvecs[0].shape[0], len(qvecs)))
# for i in range(len(qvecs)):
# qvecs_temp[:, i] = qvecs[i][:, 0].numpy()
# qvecs = qvecs_temp
#
# scores = np.zeros((len(vecs), qvecs.shape[-1]))
# for i in range(len(vecs)):
# scores[i, :] = np.amax(np.dot(vecs[i].T, qvecs), 0)
ranks = np.argsort(-scores, axis=0)
mismatched_info = compute_map_and_print(dataset, ranks, cfg['gnd'], kappas=[1, 5, 10, 100])
# hxq added
show_false_img = False
if show_false_img == True:
print('>> Save mismatched image tuple...')
for info in mismatched_info:
mismatched_img_show_save(info, qimages, images, args, bbxs=bbxs)
if Lw is not None:
# whiten the vectors
vecs_lw = whitenapply(vecs, Lw['m'], Lw['P'])
qvecs_lw = whitenapply(qvecs, Lw['m'], Lw['P'])
# search, rank, and print
scores = np.dot(vecs_lw.T, qvecs_lw)
ranks = np.argsort(-scores, axis=0)
mismatched_info = compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd'])
# hxq added
# show_false_img = False
if show_false_img == True:
print('>> Save mismatched image tuple...')
for info in mismatched_info:
mismatched_img_show_save(info, qimages, images, args, bbxs=bbxs)
print('>> {}: elapsed time: {}'.format(dataset, htime(time.time()-start)))
def main():
global args, min_loss
args = parser.parse_args()
print(args)
# set cuda visible device
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
# check if test dataset are downloaded
# and download if they are not
download_train(get_data_root())
download_test(get_data_root())
# create export dir if it doesnt exist
directory = "{}".format(args.target.replace("/", "_"))
directory += "_{}".format(args.arch)
directory += "_{}".format(args.pool)
if args.whitening:
directory += "_whiten"
if not args.pretrained:
directory += "_notpretrained"
# directory += "_bsize{}_imsize{}".format(args.batch_size, args.image_size)
directory += "_pretrained" if args.is_pretrained else ""
directory += "_random" if args.is_random else ""
directory += args.notion
target_net = args.target[6:].replace("_", "/")
print(target_net)
state = torch.load(target_net)
lw = state["meta"]["Lw"]["retrieval-SfM-120k"]["ss"]
args.directory = os.path.join(args.directory, directory)
print(">> Creating directory if it does not exist:\n>> '{}'".format(args.directory))
if not os.path.exists(args.directory):
os.makedirs(args.directory)
# set random seeds (maybe pass as argument)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
np.random.seed(0)
# create model
print(">> Using pre-trained model '{}'".format(args.arch))
model = init_network(
model=args.arch,
pooling=args.pool,
whitening=args.whitening,
pretrained=not args.is_random,
)
model.cuda()
target_model = init_network(
model=args.arch,
pooling=args.pool,
whitening=args.whitening,
pretrained=not args.is_random,
)
target_model.load_state_dict(state["state_dict"])
lw_m = lw["m"].copy()
lw_p = lw["P"].copy()
target_model.lw_m = nn.Parameter(torch.from_numpy(lw_m).float())
target_model.lw_p = nn.Parameter(torch.from_numpy(lw_p).float())
target_model.cuda()
lw_m = lw["m"].copy()
lw_p = lw["P"].copy()
model.lw_m = nn.Parameter(torch.from_numpy(lw_m).float())
model.lw_p = nn.Parameter(torch.from_numpy(lw_p).float())
whiten_layer = Whiten_layer(lw["P"].shape[1], lw["P"].shape[0])
model.white_layer = whiten_layer
model.cuda()
# parameters split into features and pool (no weight decay for pooling layer)
parameters = [
{"params": model.features.parameters()},
{"params": model.pool.parameters(), "lr": args.lr * 10, "weight_decay": 0},
{"params": model.white_layer.parameters(), "lr": 1e-2, "weight_decay": 5e-1},
]
if model.whiten is not None:
parameters.append({"params": model.whiten.parameters()})
# define optimizer
if args.optimizer == "sgd":
optimizer = torch.optim.SGD(
parameters, args.lr, momentum=args.momentum, weight_decay=args.weight_decay
)
elif args.optimizer == "adam":
optimizer = torch.optim.Adam(
parameters, args.lr, weight_decay=args.weight_decay
)
# define learning rate decay schedule
exp_decay = math.exp(-0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay)
# optionally resume from a checkpoint
start_epoch = 0
# Data loading code
normalize = transforms.Normalize(mean=model.meta["mean"], std=model.meta["std"])
transform = transforms.Compose([transforms.ToTensor(), normalize])
val_dataset = Distillation_dataset(
imsize=(args.image_size, args.image_size),
nnum=1,
qsize=float("Inf"),
poolsize=float("Inf"),
transform=transform,
filename=args.target,
q_percent=args.q_percent,
)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=10,
shuffle=False,
num_workers=args.workers,
pin_memory=False,
collate_fn=collate_tuples,
)
min_epoch = -1
for epoch in range(start_epoch, args.epochs):
if args.is_pretrained or args.is_random:
break
# set manual seeds per epoch
np.random.seed(epoch)
torch.manual_seed(epoch)
torch.cuda.manual_seed_all(epoch)
# adjust learning rate for each epoch
scheduler.step()
loss = train(val_loader, model, optimizer, epoch, target_model)
print(loss)
# evaluate on test datasets
if (epoch + 1) % 1 == 0:
with torch.no_grad():
test(args.test_datasets, model, lw)
# remember best loss and save checkpoint
is_best = loss < min_loss
min_loss = min(loss, min_loss)
save_checkpoint(
{
"epoch": epoch + 1,
"meta": model.meta,
"state_dict": model.state_dict(),
"min_loss": min_loss,
"optimizer": optimizer.state_dict(),
},
is_best,
args.directory,
)
if is_best:
min_epoch = epoch
# if epoch - min_epoch > 5:
# # break
# if val_dataset.phase == 1:
# print(bcolors.str(">>> phase 2", bcolors.OKGREEN))
# val_dataset.phase = 2
# min_epoch = epoch
# for group in optimizer.param_groups:
# group["lr"] /= 10
# else:
# break
if args.is_pretrained or args.is_random:
save_checkpoint(
{
"epoch": 0 + 1,
"meta": model.meta,
"state_dict": model.state_dict(),
"min_loss": min_loss,
"optimizer": optimizer.state_dict(),
},
True,
args.directory,
)
def main():
args = parser.parse_args()
# check if test dataset are downloaded
# and download if they are not
#download_train(get_data_root())
#download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path],
model_dir=os.path.join(get_data_root(),
'networks'))
else:
state = torch.load(args.network_path)
net = init_network(model=state['meta']['architecture'],
pooling=state['meta']['pooling'],
whitening=state['meta']['whitening'],
mean=state['meta']['mean'],
std=state['meta']['std'],
pretrained=False)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
offtheshelf = args.network_offtheshelf.split('-')
if len(offtheshelf) == 3:
if offtheshelf[2] == 'whiten':
offtheshelf_whiten = True
else:
raise (RuntimeError(
"Incorrect format of the off-the-shelf network. Examples: resnet101-gem | resnet101-gem-whiten"
))
else:
offtheshelf_whiten = False
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(model=offtheshelf[0],
pooling=offtheshelf[1],
whitening=offtheshelf_whiten)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = [1]
msp = 1
if args.multiscale:
ms = [1, 1. / math.sqrt(2), 1. / 2]
if net.meta['pooling'] == 'gem' and net.whiten is None:
msp = net.pool.p.data.tolist()[0]
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(mean=net.meta['mean'],
std=net.meta['std'])
transform = transforms.Compose([transforms.ToTensor(), normalize])
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
if args.multiscale:
Lw = net.meta['Lw'][args.whitening]['ms']
else:
Lw = net.meta['Lw'][args.whitening]['ss']
else:
print('>> {}: Learning whitening...'.format(args.whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train', args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root,
'{}-whiten.pkl'.format(args.whitening))
with open(db_fn, 'rb') as f:
db = pickle.load(f)
images = [
cid2filename(db['cids'][i], ims_root)
for i in range(len(db['cids']))
]
# extract whitening vectors
print('>> {}: Extracting...'.format(args.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
print('>> {}: elapsed time: {}'.format(args.whitening,
htime(time.time() - start)))
else:
Lw = None
# evaluate on test datasets
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
print('>> {}: Extracting...'.format(dataset))
# prepare config structure for the test dataset
cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])]
qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])]
bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net,
qimages,
args.image_size,
transform,
bbxs=bbxs,
ms=ms,
msp=msp)
print('>> {}: Evaluating...'.format(dataset))
# convert to numpy
vecs = vecs.numpy()
qvecs = qvecs.numpy()
# search, rank, and print
scores = np.dot(vecs.T, qvecs)
ranks = np.argsort(-scores, axis=0)
compute_map_and_print(dataset, ranks, cfg['gnd'])
if Lw is not None:
# whiten the vectors
vecs_lw = whitenapply(vecs, Lw['m'], Lw['P'])
qvecs_lw = whitenapply(qvecs, Lw['m'], Lw['P'])
# search, rank, and print
scores = np.dot(vecs_lw.T, qvecs_lw)
ranks = np.argsort(-scores, axis=0)
compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd'])
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
def main():
args = parser.parse_args()
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
result_dir = 'retreival_results'
if args.network_path is not None:
result_dir = os.path.join(result_dir, args.network_path)
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path],
model_dir=os.path.join(get_data_root(),
'networks'))
else:
state = torch.load(args.network_path)
net = init_network(model=state['meta']['architecture'],
pooling=state['meta']['pooling'],
whitening=state['meta']['whitening'],
mean=state['meta']['mean'],
std=state['meta']['std'],
pretrained=False)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
result_dir = os.path.join(result_dir, args.network_offtheshelf)
offtheshelf = args.network_offtheshelf.split('-')
if len(offtheshelf) == 3:
if offtheshelf[2] == 'whiten':
offtheshelf_whiten = True
else:
raise (RuntimeError(
"Incorrect format of the off-the-shelf network. Examples: resnet101-gem | resnet101-gem-whiten"
))
else:
offtheshelf_whiten = False
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(model=offtheshelf[0],
pooling=offtheshelf[1],
whitening=offtheshelf_whiten)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = [1]
msp = 1
if args.multiscale:
ms = [1, 1. / math.sqrt(2), 1. / 2]
if net.meta['pooling'] == 'gem' and net.whiten is None:
msp = net.pool.p.data.tolist()[0]
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(mean=net.meta['mean'],
std=net.meta['std'])
transform = transforms.Compose([transforms.ToTensor(), normalize])
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
if args.multiscale:
Lw = net.meta['Lw'][args.whitening]['ms']
else:
Lw = net.meta['Lw'][args.whitening]['ss']
else:
# Save whitening TODO
print('>> {}: Learning whitening...'.format(args.whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train', args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root,
'{}-whiten.pkl'.format(args.whitening))
with open(db_fn, 'rb') as f:
db = pickle.load(f)
images = [
cid2filename(db['cids'][i], ims_root)
for i in range(len(db['cids']))
]
# extract whitening vectors
print('>> {}: Extracting...'.format(args.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
print('>> {}: elapsed time: {}'.format(args.whitening,
htime(time.time() - start)))
else:
Lw = None
# evaluate on test datasets
data_root = args.data_root
datasets = datasets_names[args.dataset]
result_dict = {}
for dataset in datasets:
start = time.time()
result_dict[dataset] = {}
print('>> {}: Extracting...'.format(dataset))
# prepare config structure for the test dataset
images = get_imlist(data_root, dataset, args.train_txt)
qimages = get_imlist(data_root, dataset, args.query_txt)
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net,
images,
args.image_size,
transform,
root=os.path.join(data_root, dataset),
ms=ms,
msp=msp)
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net,
qimages,
args.image_size,
transform,
root=os.path.join(data_root, dataset),
ms=ms,
msp=msp)
print('>> {}: Evaluating...'.format(dataset))
# convert to numpy
vecs = vecs.numpy()
qvecs = qvecs.numpy()
scores, ranks = cal_ranks(vecs, vecs, Lw)
result_dict[dataset]['train'] = {'scores': scores, 'ranks': ranks}
scores, ranks = cal_ranks(vecs, qvecs, Lw)
result_dict[dataset]['test'] = {'scores': scores, 'ranks': ranks}
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
# Save retrieval results
if not os.path.exists(result_dir):
os.makedirs(result_dir)
result_file = os.path.join(result_dir, args.outfile)
np.save(result_file, result_dict)
print('Save retrieval results to {}'.format(result_file))
def main():
args = parser.parse_args()
# check if there are unknown datasets
for scene in args.scenes.split(','):
if scene not in datasets_names:
raise ValueError('Unsupported or unknown scene: {}!'.format(scene))
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network
# pretrained networks (downloaded automatically)
print(">> Loading network:\n>>>> '{}'".format(args.network))
state = load_url(PRETRAINED[args.network],
model_dir=os.path.join(get_data_root(), 'networks'))
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get('local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
# network initialization
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = list(eval(args.multiscale))
print(">>>> Evaluating scales: {}".format(ms))
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(mean=net.meta['mean'],
std=net.meta['std'])
transform = transforms.Compose([transforms.ToTensor(), normalize])
# evaluate on test datasets
scenes = args.scenes.split(',')
for scene in scenes:
start = time.time()
print('>> {}: Extracting...'.format(scene))
img_path = osp.join(args.data_path, scene, "images")
images = [
osp.join(img_path, fname) for fname in os.listdir(img_path)
if fname[-3:].lower() in ['jpg', 'png']
]
# extract vectors
vecs = extract_vectors(net, images, args.image_size, transform, ms=ms)
print('>> {}: Evaluating...'.format(scene))
# convert to numpy
vecs = vecs.numpy()
# search, rank, and print
scores = np.dot(vecs.T, vecs)
ranks = np.argsort(-scores, axis=0)
images = [img.split('/')[-1] for img in images]
for top_k in list(eval(args.top_n)):
pairs = []
for q_id in range(len(images)):
img_q = images[q_id]
pairs_per_q = [
" ".join([img_q, images[db_id]])
for db_id in list(ranks[1:top_k + 1, q_id])
]
pairs += pairs_per_q
with open(
osp.join(args.data_path, scene,
"image_pairs_" + str(top_k) + ".txt"), "w") as f:
for pair in pairs:
f.write(pair + "\n")
print('>> {}: elapsed time: {}'.format(scene,
htime(time.time() - start)))
def main():
global args, min_loss
args = parser.parse_args()
# check if test dataset are downloaded
# and download if they are not
download_train(get_data_root())
download_test(get_data_root())
# create export dir if it doesnt exist
directory = "{}".format(args.training_dataset)
directory += "_{}".format(args.arch)
directory += "_{}".format(args.pool)
if args.whitening:
directory += "_whiten"
if not args.pretrained:
directory += "_notpretrained"
directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin)
directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr,
args.weight_decay)
directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num,
args.query_size,
args.pool_size)
directory += "_bsize{}_imsize{}".format(args.batch_size, args.image_size)
args.directory = os.path.join(args.directory, directory)
print(">> Creating directory if it does not exist:\n>> '{}'".format(
args.directory))
if not os.path.exists(args.directory):
os.makedirs(args.directory)
# set cuda visible device
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# set random seeds (maybe pass as argument)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
np.random.seed(0)
# create model
if args.pretrained:
print(">> Using pre-trained model '{}'".format(args.arch))
model = init_network(model=args.arch,
pooling=args.pool,
whitening=args.whitening)
else:
print(">> Using model from scratch (random weights) '{}'".format(
args.arch))
model = init_network(model=args.arch,
pooling=args.pool,
whitening=args.whitening,
pretrained=False)
# move network to gpu
model.cuda()
# define loss function (criterion) and optimizer
if args.loss == 'contrastive':
criterion = ContrastiveLoss(margin=args.loss_margin).cuda()
else:
raise (RuntimeError("Loss {} not available!".format(args.loss)))
# parameters split into features and pool (no weight decay for pooling layer)
parameters = [{
'params': model.features.parameters()
}, {
'params': model.pool.parameters(),
'lr': args.lr * 10,
'weight_decay': 0
}]
if model.whiten is not None:
parameters.append({'params': model.whiten.parameters()})
# define optimizer
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(parameters,
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(parameters,
args.lr,
weight_decay=args.weight_decay)
# define learning rate decay schedule
exp_decay = math.exp(-0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer,
gamma=exp_decay)
# optionally resume from a checkpoint
start_epoch = 0
if args.resume:
args.resume = os.path.join(args.directory, args.resume)
if os.path.isfile(args.resume):
print(">> Loading checkpoint:\n>> '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
min_loss = checkpoint['min_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
scheduler = torch.optim.lr_scheduler.ExponentialLR(
optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch'] - 1)
else:
print(">> No checkpoint found at '{}'".format(args.resume))
# Data loading code
normalize = transforms.Normalize(mean=model.meta['mean'],
std=model.meta['std'])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
train_dataset = TuplesDataset(name=args.training_dataset,
mode='train',
imsize=args.image_size,
nnum=args.neg_num,
qsize=args.query_size,
poolsize=args.pool_size,
transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
sampler=None,
drop_last=True,
collate_fn=collate_tuples)
if args.val:
val_dataset = TuplesDataset(name=args.training_dataset,
mode='val',
imsize=args.image_size,
nnum=args.neg_num,
qsize=float('Inf'),
poolsize=float('Inf'),
transform=transform)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
drop_last=True,
collate_fn=collate_tuples)
# evaluate the network before starting
test(args.test_datasets, model)
for epoch in range(start_epoch, args.epochs):
# set manual seeds per epoch
np.random.seed(epoch)
torch.manual_seed(epoch)
torch.cuda.manual_seed_all(epoch)
# adjust learning rate for each epoch
scheduler.step()
# lr_feat = optimizer.param_groups[0]['lr']
# lr_pool = optimizer.param_groups[1]['lr']
# print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool))
# train for one epoch on train set
loss = train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
if args.val:
loss = validate(val_loader, model, criterion, epoch)
# evaluate on test datasets
test(args.test_datasets, model)
# remember best loss and save checkpoint
is_best = loss < min_loss
min_loss = min(loss, min_loss)
save_checkpoint(
{
'epoch': epoch + 1,
'meta': model.meta,
'state_dict': model.state_dict(),
'min_loss': min_loss,
'optimizer': optimizer.state_dict(),
}, is_best, args.directory)
def main():
global args, min_loss
args = parser.parse_args()
# set cuda visible device
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
# check if test dataset are downloaded
# and download if they are not
download_train(get_data_root())
download_test(get_data_root())
# set random seeds (maybe pass as argument)
torch.manual_seed(1234)
torch.cuda.manual_seed_all(1234)
np.random.seed(1234)
random.seed(1234)
torch.backends.cudnn.deterministic = True
state = torch.load(args.network_path)
model = init_network(
model=state["meta"]["architecture"],
pooling=state["meta"]["pooling"],
whitening=state["meta"]["whitening"],
mean=state["meta"]["mean"],
std=state["meta"]["std"],
pretrained=False,
)
model.load_state_dict(state["state_dict"])
model.meta["Lw"] = state["meta"]["Lw"]
model.cuda()
# whitening
Lw = model.meta["Lw"]["retrieval-SfM-120k"]["ss"]
Lw_m = torch.from_numpy(Lw["m"]).cuda().float()
Lw_p = torch.from_numpy(Lw["P"]).cuda().float()
# Data loading code
normalize = transforms.Normalize(mean=model.meta["mean"],
std=model.meta["std"])
transform = transforms.Compose([transforms.ToTensor(), normalize])
query_dataset = MyTripletDataset(
# imsize=args.image_size,
imsize=(362, 362),
random=False,
transform=transform,
norm=(0, 0),
filename="base/" + args.network_path.replace("/", "_") + "_triplet",
)
# val_dataset.test_cluster(model)
# return
query_dataset.create_epoch_tuples(model)
query_loader = torch.utils.data.DataLoader(
query_dataset,
batch_size=1,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
worker_init_fn=lambda _: random.seed(1234),
)
base_dataset = ImagesFromList(
root="",
images=query_dataset.images,
# imsize=query_dataset.imsize,
imsize=(362, 362),
transform=query_dataset.transform,
random=False,
)
base_loader = torch.utils.data.DataLoader(base_dataset,
batch_size=1,
shuffle=False)
# test(["oxford5k"], model)
extract(query_loader, base_loader, model, Lw_m, Lw_p)
def main():
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('train_images_path')
arg_parser.add_argument('test_images_path')
arg_parser.add_argument('predictions_path')
args = arg_parser.parse_args()
imsize=480
train_path=args.train_images_path
test_path=args.test_images_path
outfile=args.predictions_path
##read data##
data_list=os.listdir(ipath)
train_images = get_imlist(train_path)
test_images = get_imlist(test_path)
##RAMAC##
RAMAC = extract_feature(train_images, 'resnet101', imsize)
RAMAC_test = extract_feature(test_images, 'resnet101', imsize)
##UEL##
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform_train = transforms.Compose([
transforms.ToPILImage(),
transforms.RandomCrop(size=224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
transform_test = transforms.Compose([
transforms.ToPILImage(),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])
net = resnet101(pretrained=True,low_dim=128)
model_path = './model/UEL.t'#After training UEL
net.load_state_dict(torch.load())
imset = DataLoader(path = train_path, transform=transform_test)
train_loader = torch.utils.data.DataLoader(imset, batch_size=32, shuffle=False, num_workers=0)
UEL = obtainf(net, train_loader)
imset = DataLoader(path = test_path, transform=transform_test)
test_loader = torch.utils.data.DataLoader(imset, batch_size=32, shuffle=False, num_workers=0)
UEL_test = obtainf(net, test_loader)
##GEM##
image_size=1024
multiscale='[1, 2**(1/2), 1/2**(1/2)]'
state = torch.load('./model/retrievalSfM120k-vgg16-gem-b4dcdc6.pth')
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get('local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
# load network
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
ms = list(eval(multiscale))
msp = net.pool.p.item()
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(
mean=net.meta['mean'],
std=net.meta['std']
)
transform = transforms.Compose([
transforms.ToTensor(),
normalize
])
GEM = extract_vectors(net,train_images , 480, transform, ms=ms, msp=msp).numpy().T
GEM_test = extract_vectors(net,test_images , 480, transform, ms=ms, msp=msp).numpy().T
##Retrieval##
feats=np.concatenate((RAMAC,UEL,GEM),axis=1).astype('float32')
query_feat=np.concatenate((RAMAC_test, UEL_test,GEM_test),axis=1).astype('float32')
##diffusion##
kq, kd = 7, 50
gamma=80
diffusion = Diffusion(feats, '/')
offline = diffusion.get_offline_results(1024, kd)
print('[search] 1) k-NN search')
sims, ids = diffusion.knn.search(query_feat, kq)
sims = sims ** gamma
qr_num = ids.shape[0]
print('[search] 2) linear combination')
all_scores = np.empty((qr_num, 7), dtype=np.float32)
all_ranks = np.empty((qr_num, 7), dtype=np.int)
for i in range(qr_num):
scores = sims[i] @ offline[ids[i]]
parts = np.argpartition(-scores, 7)[:7]
ranks = np.argsort(-scores[parts])
all_scores[i] = scores[parts][ranks]
all_ranks[i] = parts[ranks]
I = all_ranks
##output##
out=pd.DataFrame(list(map(lambda x: x.split('/')[-1].split('.jpg')[0],timages)))
out['1']=pd.DataFrame(I)[0].map(lambda x:data_list[x].split('.')[0] )
out['2']=pd.DataFrame(I)[1].map(lambda x:data_list[x].split('.')[0] )
out['3']=pd.DataFrame(I)[2].map(lambda x:data_list[x].split('.')[0] )
out['4']=pd.DataFrame(I)[3].map(lambda x:data_list[x].split('.')[0] )
out['5']=pd.DataFrame(I)[4].map(lambda x:data_list[x].split('.')[0] )
out['6']=pd.DataFrame(I)[5].map(lambda x:data_list[x].split('.')[0] )
out['7']=pd.DataFrame(I)[6].map(lambda x:data_list[x].split('.')[0] )
out.to_csv(outfile,index=None,header=None)
N, C, H, W = 10, 3, 256, 256
D_out = 2048
x = torch.randn(N, C, H, W)
y = torch.randn(D_out, N)
model_params = {}
model_params['architecture'] = 'resnet101'
model_params['pooling'] = 'gem'
model_params['local_whitening'] = False
model_params['regional'] = False
model_params['whitening'] = False
# model_params['mean'] = ... # will use default
# model_params['std'] = ... # will use default
model_params['pretrained'] = True
model = init_network(model_params)
# model = torch.nn.Sequential(
# DataParallelModel(D_in, H),
# DataParallelModel(H, H),
# torch.nn.Linear(H, D_out)
# )
model.cuda()
criterion = ContrastiveLoss(margin=0.85).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# # Data loading code
# normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std'])
# transform = transforms.Compose([
def main():
global args, min_loss
args = parser.parse_args()
# create export dir if it doesnt exist
conftxt = 'Model {}-{}'.format(args.arch, args.pool)
if args.whitening:
conftxt += '-whiten'
if args.pretrained in PRETRAINED:
conftxt += " Pretrained: {} ".format(args.pretrained)
elif args.pretrained == 'imagenet':
conftxt += " Pretrained: imagenet"
else:
conftxt += " No pretrain "
conftxt += " Loss: {} margin: {:.2f}".format(args.loss, args.loss_margin)
conftxt += " Opt: {} lr{:.1e} wd{:.1e}".format(args.optimizer, args.lr,
args.weight_decay)
conftxt += " snn{} dnn{}".format(args.sneg_num, args.dneg_num)
conftxt += " qsize{}".format(args.query_size)
conftxt += " bsize{} imsize{}".format(args.batch_size, args.image_size)
print(">> Creating directory if it does not exist:\n>> '{}'".format(
args.directory))
if not os.path.exists(args.directory):
os.makedirs(args.directory)
log = open(os.path.join(args.directory, 'log.txt'), 'w')
lprint('>>>>Training configuration:\n {}'.format(conftxt), log)
# set cuda visible device
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# set random seeds (maybe pass as argument)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
np.random.seed(0)
# create model
if args.pretrained:
if args.pretrained in PRETRAINED:
# pretrained networks (downloaded automatically)
print(">> Load pre-trained model from '{}'".format(
PRETRAINED[args.pretrained]))
state = load_url(PRETRAINED[args.pretrained],
model_dir=os.path.join('data/networks'))
model = init_network(model=state['meta']['architecture'],
pooling=state['meta']['pooling'],
whitening=state['meta']['whitening'],
mean=state['meta']['mean'],
std=state['meta']['std'],
pretrained=False)
model.load_state_dict(state['state_dict'])
else:
print(">> Using pre-trained model on imagenet '{}'".format(
args.arch))
model = init_network(model=args.arch,
pooling=args.pool,
whitening=args.whitening,
pretrained=True)
else:
print(">> Using model from scratch (random weights) '{}'".format(
args.arch))
model = init_network(model=args.arch,
pooling=args.pool,
whitening=args.whitening,
pretrained=False)
# move network to gpu
model.cuda()
# define loss function (criterion) and optimizer
if args.loss == 'contrastive':
criterion = ContrastiveLoss(margin=args.loss_margin).cuda()
else:
raise (RuntimeError("Loss {} not available!".format(args.loss)))
# parameters split into features and pool (no weight decay for pooling layer)
parameters = [{
'params': model.features.parameters()
}, {
'params': model.pool.parameters(),
'lr': args.lr * 10,
'weight_decay': 0
}]
if model.whiten is not None:
parameters.append({'params': model.whiten.parameters()})
# define optimizer
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(parameters,
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(parameters,
args.lr,
weight_decay=args.weight_decay)
# define learning rate decay schedule
exp_decay = math.exp(-0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer,
gamma=exp_decay)
# optionally resume from a checkpoint
start_epoch = 0
if args.resume:
args.resume = os.path.join(args.directory, args.resume)
if os.path.isfile(args.resume):
print(">> Loading checkpoint:\n>> '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
min_loss = checkpoint['min_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
scheduler = torch.optim.lr_scheduler.ExponentialLR(
optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch'] - 1)
else:
print(">> No checkpoint found at '{}'".format(args.resume))
# Data loading code
normalize = transforms.Normalize(mean=model.meta['mean'],
std=model.meta['std'])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
train_dataset = CustomizeTuplesDataset(name='7Scene',
mode='train',
db_file=args.db_file,
ims_root=args.data_root,
imsize=args.image_size,
snum=args.sneg_num,
dnum=args.dneg_num,
qsize=args.query_size,
transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
sampler=None,
drop_last=True,
collate_fn=collate_tuples)
if args.val:
val_dataset = CustomizeTuplesDataset(name='7Scene',
mode='val',
db_file=args.db_file,
ims_root=args.data_root,
imsize=args.image_size,
snum=args.sneg_num,
dnum=args.dneg_num,
qsize=float('Inf'),
transform=transform)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
drop_last=True,
collate_fn=collate_tuples)
# evaluate the network before starting
gt_root = os.path.join(args.data_root, 'sfm_relative_pose_pairs')
data_splits = split_dataset(args.data_root,
test_datasets,
val_step=6,
seed=0)
test(model,
args.data_root,
data_splits,
gt_root,
epoch=0,
pass_thres=8,
knn=10,
query_key='val',
db_key='train',
log=log)
for epoch in range(start_epoch, args.epochs):
# set manual seeds per epoch
np.random.seed(epoch)
torch.manual_seed(epoch)
torch.cuda.manual_seed_all(epoch)
# adjust learning rate for each epoch
scheduler.step()
# train for one epoch on train set
loss = train(train_loader,
model,
criterion,
optimizer,
epoch,
print_freq=40,
log=log)
# evaluate on validation set
if args.val and (epoch + 1) % 5 == 0:
loss = validate(val_loader,
model,
criterion,
epoch,
print_freq=100,
log=log)
# evaluate on test datasets
test(model,
args.data_root,
data_splits,
gt_root,
epoch,
pass_thres=8,
knn=10,
query_key='val',
db_key='train',
log=log)
# remember best loss and save checkpoint
is_best = loss < min_loss
min_loss = min(loss, min_loss)
save_checkpoint(
{
'epoch': epoch + 1,
'meta': model.meta,
'state_dict': model.state_dict(),
'min_loss': min_loss,
'optimizer': optimizer.state_dict(),
}, is_best, args.directory)
print('Training Finished')
log.close()
def main():
args = parser.parse_args()
# check if there are unknown datasets
for dataset in args.datasets.split(','):
if dataset not in datasets_names:
raise ValueError(
'Unsupported or unknown dataset: {}!'.format(dataset))
# check if test dataset are downloaded
# and download if they are not
download_train(get_data_root())
download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path],
model_dir=os.path.join(get_data_root(),
'networks'))
else:
# fine-tuned network from path
state = torch.load(args.network_path)
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get(
'local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
# load network
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
# parse off-the-shelf parameters
offtheshelf = args.network_offtheshelf.split('-')
net_params = {}
net_params['architecture'] = offtheshelf[0]
net_params['pooling'] = offtheshelf[1]
net_params['local_whitening'] = 'lwhiten' in offtheshelf[2:]
net_params['regional'] = 'reg' in offtheshelf[2:]
net_params['whitening'] = 'whiten' in offtheshelf[2:]
net_params['pretrained'] = True
# load off-the-shelf network
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(net_params)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = list(eval(args.multiscale))
if len(ms) > 1 and net.meta['pooling'] == 'gem' and not net.meta[
'regional'] and not net.meta['whitening']:
msp = net.pool.p.item()
print(">> Set-up multiscale:")
print(">>>> ms: {}".format(ms))
print(">>>> msp: {}".format(msp))
else:
msp = 1
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(mean=net.meta['mean'],
std=net.meta['std'])
transform = transforms.Compose([transforms.ToTensor(), normalize])
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
if len(ms) > 1:
Lw = net.meta['Lw'][args.whitening]['ms']
else:
Lw = net.meta['Lw'][args.whitening]['ss']
else:
# if we evaluate networks from path we should save/load whitening
# not to compute it every time
if args.network_path is not None:
whiten_fn = args.network_path + '_{}_whiten'.format(
args.whitening)
if len(ms) > 1:
whiten_fn += '_ms'
whiten_fn += '.pth'
else:
whiten_fn = None
if whiten_fn is not None and os.path.isfile(whiten_fn):
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
Lw = torch.load(whiten_fn)
else:
print('>> {}: Learning whitening...'.format(args.whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train',
args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root,
'{}-whiten.pkl'.format(args.whitening))
with open(db_fn, 'rb') as f:
db = pickle.load(f)
images = [
cid2filename(db['cids'][i], ims_root)
for i in range(len(db['cids']))
]
# extract whitening vectors
print('>> {}: Extracting...'.format(args.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
# saving whitening if whiten_fn exists
if whiten_fn is not None:
print('>> {}: Saving to {}...'.format(
args.whitening, whiten_fn))
torch.save(Lw, whiten_fn)
print('>> {}: elapsed time: {}'.format(args.whitening,
htime(time.time() - start)))
else:
Lw = None
# evaluate on test datasets
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
print('>> {}: Extracting...'.format(dataset))
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
images = get_imlist("E:\\PycharmProjects\\image-retrieval\\holiday2\\")
names = []
for i, img_path in enumerate(images):
img_name = os.path.split(img_path)[1]
print(img_name)
names.append(img_name)
# prepare config structure for the test dataset
# cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
# images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])]
# try:
# bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
# except:
# bbxs = None # for holidaysmanrot and copydays
# names = []
# for i, img_path in enumerate(images):
# img_name = os.path.split(img_path)[1]
# print(img_name)
# names.append(img_name)
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# convert to numpy
vecs = vecs.numpy()
vecs = vecs.T
print("--------------------------------------------------")
print(" writing feature extraction results ...")
print("--------------------------------------------------")
output = "gem_res_holiday_3.h5"
h5f = h5py.File(output, 'w')
h5f.create_dataset('dataset_1', data=vecs)
h5f.create_dataset('dataset_2', data=np.string_(names))
h5f.close()
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
state = torch.load(network_path)
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get('local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
net_params['multi_layer_cat'] = state['meta']['multi_layer_cat']
# load network
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif network_offtheshelf is not None:
# parse off-the-shelf parameters
offtheshelf = network_offtheshelf.split('-')
net_params = {}
net_params['architecture'] = offtheshelf[0]
def main():
global args, min_loss, global_step
args = parser.parse_args()
if args.wandb:
# initialize wandb
wandb.init(
project='cnnimageretrieval-pytorch', name=args.directory,
entity='ronaldseoh')
# save args provided for this experiment to wandb
wandb.config.update(args)
# manually check if there are unknown test datasets
for dataset in args.test_datasets.split(','):
if dataset not in test_datasets_names:
raise ValueError('Unsupported or unknown test dataset: {}!'.format(dataset))
# check if test dataset are downloaded
# and download if they are not
download_train(get_data_root())
download_test(get_data_root())
# create export dir if it doesnt exist
directory = "{}".format(args.training_dataset)
directory += "_{}".format(args.arch)
directory += "_{}".format(args.pool)
if args.local_whitening:
directory += "_lwhiten"
if args.regional:
directory += "_r"
if args.whitening:
directory += "_whiten"
if not args.pretrained:
directory += "_notpretrained"
directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin)
directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr, args.weight_decay)
directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num, args.query_size, args.pool_size)
directory += "_bsize{}_uevery{}_imsize{}".format(args.batch_size, args.update_every, args.image_size)
args.directory = os.path.join(args.directory, directory)
print(">> Creating directory if it does not exist:\n>> '{}'".format(args.directory))
if not os.path.exists(args.directory):
os.makedirs(args.directory)
# set cuda visible device
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# set random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
# initialize model
if args.pretrained:
print(">> Using pre-trained model '{}'".format(args.arch))
else:
print(">> Using model from scratch (random weights) '{}'".format(args.arch))
model_params = {}
model_params['architecture'] = args.arch
model_params['pooling'] = args.pool
model_params['local_whitening'] = args.local_whitening
model_params['regional'] = args.regional
model_params['whitening'] = args.whitening
# model_params['mean'] = ... # will use default
# model_params['std'] = ... # will use default
model_params['pretrained'] = args.pretrained
model = init_network(model_params)
# move network to gpu
model.cuda()
# define loss function (criterion) and optimizer
if args.loss == 'contrastive':
criterion = ContrastiveLoss(margin=args.loss_margin).cuda()
elif args.loss == 'triplet':
criterion = TripletLoss(margin=args.loss_margin).cuda()
else:
raise(RuntimeError("Loss {} not available!".format(args.loss)))
# parameters split into features, pool, whitening
# IMPORTANT: no weight decay for pooling parameter p in GeM or regional-GeM
parameters = []
# add feature parameters
parameters.append({'params': model.features.parameters()})
# add local whitening if exists
if model.lwhiten is not None:
parameters.append({'params': model.lwhiten.parameters()})
# add pooling parameters (or regional whitening which is part of the pooling layer!)
if not args.regional:
# global, only pooling parameter p weight decay should be 0
if args.pool == 'gem':
parameters.append({'params': model.pool.parameters(), 'lr': args.lr*10, 'weight_decay': 0})
elif args.pool == 'gemmp':
parameters.append({'params': model.pool.parameters(), 'lr': args.lr*100, 'weight_decay': 0})
else:
# regional, pooling parameter p weight decay should be 0,
# and we want to add regional whitening if it is there
if args.pool == 'gem':
parameters.append({'params': model.pool.rpool.parameters(), 'lr': args.lr*10, 'weight_decay': 0})
elif args.pool == 'gemmp':
parameters.append({'params': model.pool.rpool.parameters(), 'lr': args.lr*100, 'weight_decay': 0})
if model.pool.whiten is not None:
parameters.append({'params': model.pool.whiten.parameters()})
# add final whitening if exists
if model.whiten is not None:
parameters.append({'params': model.whiten.parameters()})
# define optimizer
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(parameters, args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(parameters, args.lr, weight_decay=args.weight_decay)
# define learning rate decay schedule
# TODO: maybe pass as argument in future implementation?
exp_decay = math.exp(-0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay)
# optionally resume from a checkpoint
start_epoch = 0
if args.resume:
args.resume = os.path.join(args.directory, args.resume)
if os.path.isfile(args.resume):
# load checkpoint weights and update model and optimizer
print(">> Loading checkpoint:\n>> '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
min_loss = checkpoint['min_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
# important not to forget scheduler updating
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch']-1)
else:
print(">> No checkpoint found at '{}'".format(args.resume))
# Data loading code
normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std'])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
if args.query_size == -1:
train_query_size = float('Inf')
else:
train_query_size = args.query_size
if args.pool_size == -1:
train_pool_size = float('Inf')
else:
train_pool_size = args.pool_size
train_dataset = TuplesDataset(
name=args.training_dataset,
mode='train',
imsize=args.image_size,
nnum=args.neg_num,
qsize=train_query_size,
poolsize=train_pool_size,
transform=transform,
store_nidxs_others_up_to=args.store_nidxs_others_up_to,
store_nidxs_others_order_by=args.store_nidxs_others_order_by,
totally_random_nidxs=args.totally_random_nidxs,
totally_random_nidxs_others=args.totally_random_nidxs_others,
dense_refresh_batch_and_nearby=args.dense_refresh_batch_and_nearby,
dense_refresh_batch_multi_hop=args.dense_refresh_batch_multi_hop,
dense_refresh_batch_random=args.dense_refresh_batch_random,
)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=0, pin_memory=True, sampler=None,
drop_last=True, collate_fn=collate_tuples
)
if args.val:
val_dataset = TuplesDataset(
name=args.training_dataset,
mode='val',
imsize=args.image_size,
nnum=args.neg_num,
qsize=float('Inf'),
poolsize=float('Inf'),
transform=transform
)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True,
drop_last=True, collate_fn=collate_tuples
)
if args.wandb:
# Start watching 'model' from wandb
wandb.watch(model)
# evaluate the network before starting
# this might not be necessary?
test(args.test_datasets, model, wandb_enabled=args.wandb, epoch=-1)
indexes_to_refresh = []
for epoch in range(start_epoch, args.epochs):
# set manual seeds per epoch
random.seed(args.seed + epoch)
np.random.seed(args.seed + epoch)
torch.manual_seed(args.seed + epoch)
torch.cuda.manual_seed_all(args.seed + epoch)
# train for one epoch on train set
loss, indexes_to_refresh = train(train_loader, model, criterion, optimizer, epoch, indexes_to_refresh)
if args.wandb:
wandb.log({"loss_avg": loss, "epoch": epoch, "global_step": global_step}) ## This is average loss
# adjust learning rate for each epoch
scheduler.step()
# # debug printing to check if everything ok
# lr_feat = optimizer.param_groups[0]['lr']
# lr_pool = optimizer.param_groups[1]['lr']
# print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool))
# evaluate on validation set
if args.val:
with torch.no_grad():
loss = validate(val_loader, model, criterion, epoch)
if args.wandb:
wandb.log({"loss_validation": loss, "epoch": epoch, "global_step": global_step})
# evaluate on test datasets every test_freq epochs
if (epoch + 1) % args.test_freq == 0:
with torch.no_grad():
test(args.test_datasets, model, wandb_enabled=args.wandb, epoch=epoch)
# remember best loss and save checkpoint
is_best = loss < min_loss
min_loss = min(loss, min_loss)
if is_best:
print("Epoch", str(epoch + 1), "lower loss:", min_loss)
save_checkpoint({
'epoch': epoch + 1,
'meta': model.meta,
'state_dict': model.state_dict(),
'min_loss': min_loss,
'optimizer' : optimizer.state_dict(),
}, is_best, args.directory)
# calculate avg_neg_distance and avg_pos_distance for one last time
print("Training finished. Calculating the final avg_{neg,pos}_distance...")
avg_neg_distance, _ = train_loader.dataset.create_epoch_tuples(
model,
batch_members=[],
refresh_query_selection=False,
refresh_query_vectors=False,
refresh_negative_pool=False,
refresh_negative_pool_vectors=False,
refresh_nidxs=False,
refresh_nidxs_vectors=False)
if args.wandb:
wandb.log({"avg_neg_distance": avg_neg_distance, 'epoch': epoch, "global_step": global_step})
if args.calculate_positive_distance:
avg_pos_distance = train_loader.dataset.calculate_average_positive_distance()
if args.wandb:
wandb.log({"avg_pos_distance": avg_pos_distance, 'epoch': epoch, "global_step": global_step})
def main():
args = parser.parse_args()
# check if there are unknown datasets
for dataset in args.datasets.split(','):
if dataset not in datasets_names:
raise ValueError(
'Unsupported or unknown dataset: {}!'.format(dataset))
# check if test dataset are downloaded
# and download if they are not
#download_train(get_data_root())
download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path],
model_dir=os.path.join(get_data_root(),
'networks'))
else:
# fine-tuned network from path
state = torch.load(args.network_path)
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get(
'local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
# load network
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
# parse off-the-shelf parameters
offtheshelf = args.network_offtheshelf.split('-')
net_params = {}
net_params['architecture'] = offtheshelf[0]
net_params['pooling'] = offtheshelf[1]
net_params['local_whitening'] = 'lwhiten' in offtheshelf[2:]
net_params['regional'] = 'reg' in offtheshelf[2:]
net_params['whitening'] = 'whiten' in offtheshelf[2:]
net_params['pretrained'] = True
# load off-the-shelf network
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(net_params)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = list(eval(args.multiscale))
if len(ms) > 1 and net.meta['pooling'] == 'gem' and not net.meta[
'regional'] and not net.meta['whitening']:
msp = net.pool.p.item()
print(">> Set-up multiscale:")
print(">>>> ms: {}".format(ms))
print(">>>> msp: {}".format(msp))
else:
msp = 1
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(mean=net.meta['mean'],
std=net.meta['std'])
transform = transforms.Compose([transforms.ToTensor(), normalize])
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
if len(ms) > 1:
Lw = net.meta['Lw'][args.whitening]['ms']
else:
Lw = net.meta['Lw'][args.whitening]['ss']
else:
# if we evaluate networks from path we should save/load whitening
# not to compute it every time
if args.network_path is not None:
whiten_fn = args.network_path + '_{}_whiten'.format(
args.whitening)
if len(ms) > 1:
whiten_fn += '_ms'
whiten_fn += '.pth'
else:
whiten_fn = None
if whiten_fn is not None and os.path.isfile(whiten_fn):
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
Lw = torch.load(whiten_fn)
else:
print('>> {}: Learning whitening...'.format(args.whitening))
# loading db
db_root = os.path.join(get_data_root(), 'train',
args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root,
'{}-whiten.pkl'.format(args.whitening))
with open(db_fn, 'rb') as f:
db = pickle.load(f)
images = [
cid2filename(db['cids'][i], ims_root)
for i in range(len(db['cids']))
]
# extract whitening vectors
print('>> {}: Extracting...'.format(args.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw = {'m': m, 'P': P}
# saving whitening if whiten_fn exists
if whiten_fn is not None:
print('>> {}: Saving to {}...'.format(
args.whitening, whiten_fn))
torch.save(Lw, whiten_fn)
print('>> {}: elapsed time: {}'.format(args.whitening,
htime(time.time() - start)))
else:
Lw = None
# evaluate on test datasets
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
"""
print('>> {}: Extracting...'.format(dataset))
# prepare config structure for the test dataset
cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
images = [cfg['im_fname'](cfg,i) for i in range(cfg['n'])]
qimages = [cfg['qim_fname'](cfg,i) for i in range(cfg['nq'])]
try:
bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
except:
bbxs = None # for holidaysmanrot and copydays
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net, images, args.image_size, transform, ms=ms, msp=msp)
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net, qimages, args.image_size, transform, bbxs=bbxs, ms=ms, msp=msp)
print('>> {}: Evaluating...'.format(dataset))
# convert to numpy
vecs = vecs.numpy()
qvecs = qvecs.numpy()
print (vecs.shape)
print (qvecs.shape)
# search, rank, and print
scores = np.dot(vecs.T, qvecs)
print (scores.shape)
# to save scores (single query)
# oxford
#f = 'oxf_single.npy'
# paris
#f = 'par_single.npy'
# roxford
#f = 'roxf_single.npy'
# rparis
f = 'rpar_single.npy'
ranks = np.argsort(-scores, axis=0)
compute_map_and_print(dataset, ranks, cfg['gnd'])
if Lw is not None:
# whiten the vectors
vecs_lw = whitenapply(vecs, Lw['m'], Lw['P'])
qvecs_lw = whitenapply(qvecs, Lw['m'], Lw['P'])
# search, rank, and print
scores = np.dot(vecs_lw.T, qvecs_lw)
# save
np.save(f, scores)
ranks = np.argsort(-scores, axis=0)
compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd'])
"""
############################################################
# Test
# prepare config structure for the test dataset
cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])]
qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])]
print(qimages)
# to load scores
# oxford
#f = 'oxf_single.npy'
#f = 'oxf_mq_avg.npy'
#f = 'oxf_mq_max.npy'
#f = 'oxf_sc_imf.npy'
# paris
#f = 'par_single.npy'
#f = 'par_mq_avg.npy'
#f = 'par_mq_max.npy'
f = 'par_sc_imf.npy'
# roxford
#f = 'roxf_single.npy'
#f = 'roxf_mq_avg.npy'
#f = 'roxf_mq_max.npy'
#f = 'roxf_sc_imf.npy'
# rparis
#f = 'rpar_single.npy'
#f = 'rpar_mq_avg.npy'
#f = 'rpar_mq_max.npy'
#f = 'rpar_sc_imf.npy'
# load
scores = np.load(f)
ranks = np.argsort(-scores, axis=0)
compute_map_and_print(dataset + ' + whiten', ranks, cfg['gnd'])
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
def main():
args = parser.parse_args()
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network
# pretrained networks (downloaded automatically)
print(">> Loading network:\n>>>> '{}'".format(args.network))
state = load_url(PRETRAINED[args.network],
model_dir=os.path.join(get_data_root(), 'networks'))
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get('local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
# network initialization
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = list(eval(args.multiscale))
print(">>>> Evaluating scales: {}".format(ms))
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(mean=net.meta['mean'],
std=net.meta['std'])
transform = transforms.Compose([transforms.ToTensor(), normalize])
# evaluate on test datasets
start = time.time()
print('>> {}: Extracting...')
img_path = osp.join(args.data_path, "images", "images_upright")
fnames_db = [
osp.join(path, fname)
for path, _, files in os.walk(osp.join(img_path, "db"))
for fname in files if fname[-3:].lower() in ["jpg", "png"]
]
fnames_q = [
osp.join(path, fname)
for path, _, files in os.walk(osp.join(img_path, "query"))
for fname in files if fname[-3:].lower() in ["jpg", "png"]
]
# extract vectors
vecs_db = extract_vectors(net,
fnames_db,
args.image_size,
transform,
ms=ms)
vecs_q = extract_vectors(net, fnames_q, args.image_size, transform, ms=ms)
print('>> {}: Evaluating...')
# convert to numpy
vecs_db = vecs_db.numpy()
vecs_q = vecs_q.numpy()
# search, rank, and print
scores_db = np.dot(vecs_db.T, vecs_db)
scores_q = np.dot(vecs_db.T, vecs_q)
ranks_db = np.argsort(-scores_db, axis=0)
ranks_q = np.argsort(-scores_q, axis=0)
print(ranks_q.shape, ranks_db.shape)
images_db = [fname_db[len(img_path) + 1:] for fname_db in fnames_db]
images_q = [fname_q[len(img_path) + 1:] for fname_q in fnames_q]
pairs_db = []
for q_id in range(len(images_db)):
img_q = images_db[q_id]
pairs_per_q = [
" ".join([img_q, images_db[db_id]])
for db_id in list(ranks_db[1:args.top_n + 1, q_id])
]
pairs_db += pairs_per_q
for q_id in range(len(images_q)):
img_q = images_q[q_id]
pairs_per_q = [
" ".join([img_q, images_db[db_id]])
for db_id in list(ranks_q[1:args.top_n + 1, q_id])
]
pairs_db += pairs_per_q
with open(
osp.join(args.data_path,
"image_pairs_day_night_" + str(args.top_n) + ".txt"),
"w") as f:
for pair in pairs_db:
f.write(pair + "\n")
print("Done")
def main():
args = parser.parse_args()
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path],
model_dir=os.path.join(get_data_root(),
'networks'))
else:
# fine-tuned network from path
state = torch.load(args.network_path)
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get(
'local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
# load network
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
if "epoch" in state:
print("Model after {} epochs".format(state["epoch"]))
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
# parse off-the-shelf parameters
offtheshelf = args.network_offtheshelf.split('-')
net_params = {}
net_params['architecture'] = offtheshelf[0]
net_params['pooling'] = offtheshelf[1]
net_params['local_whitening'] = 'lwhiten' in offtheshelf[2:]
net_params['regional'] = 'reg' in offtheshelf[2:]
net_params['whitening'] = 'whiten' in offtheshelf[2:]
net_params['pretrained'] = True
# load off-the-shelf network
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(net_params)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters: test both single scale and multiscale
ms_singlescale = [1]
msp_singlescale = 1
ms_multiscale = list(eval(args.multiscale))
msp_multiscale = 1
if len(ms_multiscale
) > 1 and net.meta['pooling'] == 'gem' and not net.meta[
'regional'] and not net.meta['whitening']:
msp_multiscale = net.pool.p.item()
print(">> Set-up multiscale:")
print(">>>> ms: {}".format(ms_multiscale))
print(">>>> msp: {}".format(msp_multiscale))
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(mean=net.meta['mean'],
std=net.meta['std'])
transform = transforms.Compose([transforms.ToTensor(), normalize])
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
Lw = net.meta['Lw'][args.whitening]
else:
# if we evaluate networks from path we should save/load whitening
# not to compute it every time
if args.network_path is not None:
whiten_fn = args.network_path + '_{}_whiten'.format(
args.whitening)
whiten_fn += '.pth'
else:
whiten_fn = None
if whiten_fn is not None and os.path.isfile(whiten_fn):
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
Lw = torch.load(whiten_fn)
else:
Lw = {}
for whiten_type, ms, msp in zip(
["ss", "ms"], [ms_singlescale, ms_multiscale],
[msp_singlescale, msp_multiscale]):
print('>> {0}: Learning whitening {1}...'.format(
args.whitening, whiten_type))
# loading db
db_root = os.path.join(get_data_root(), 'train',
args.whitening)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(
db_root, '{}-whiten.pkl'.format(args.whitening))
with open(db_fn, 'rb') as f:
db = pickle.load(f)
images = [
cid2filename(db['cids'][i], ims_root)
for i in range(len(db['cids']))
]
# extract whitening vectors
print('>> {}: Extracting...'.format(args.whitening))
wvecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
wvecs = wvecs.numpy()
m, P = whitenlearn(wvecs, db['qidxs'], db['pidxs'])
Lw[whiten_type] = {'m': m, 'P': P}
print('>> {}: elapsed time: {}'.format(
args.whitening, htime(time.time() - start)))
# saving whitening if whiten_fn exists
if whiten_fn is not None:
print('>> {}: Saving to {}...'.format(
args.whitening, whiten_fn))
torch.save(Lw, whiten_fn)
else:
Lw = None
# evaluate on test datasets
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
for whiten_type, ms, msp in zip(["ss", "ms"],
[ms_singlescale, ms_multiscale],
[msp_singlescale, msp_multiscale]):
print('>> Extracting feature on {0}, whitening {1}'.format(
dataset, whiten_type))
# prepare config structure for the test dataset
cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])]
qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])]
bbxs = [tuple(cfg['gnd'][i]['bbx']) for i in range(cfg['nq'])]
# extract database and query vectors
print('>> {}: database images...'.format(dataset))
vecs = extract_vectors(net,
images,
args.image_size,
transform,
ms=ms,
msp=msp)
print('>> {}: query images...'.format(dataset))
qvecs = extract_vectors(net,
qimages,
args.image_size,
transform,
bbxs=bbxs,
ms=ms,
msp=msp)
print('>> {}: Evaluating...'.format(dataset))
# convert to numpy
vecs = vecs.numpy()
qvecs = qvecs.numpy()
# search, rank, and print
scores = np.dot(vecs.T, qvecs)
ranks = np.argsort(-scores, axis=0)
compute_map_and_print(dataset, ranks, cfg['gnd'])
if Lw is not None:
# whiten the vectors
vecs_lw = whitenapply(vecs, Lw[whiten_type]['m'],
Lw[whiten_type]['P'])
qvecs_lw = whitenapply(qvecs, Lw[whiten_type]['m'],
Lw[whiten_type]['P'])
# search, rank, and print
scores = np.dot(vecs_lw.T, qvecs_lw)
ranks = np.argsort(-scores, axis=0)
compute_map_and_print(
dataset + ' + whiten {}'.format(whiten_type), ranks,
cfg['gnd'])
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
def main():
args = parser.parse_args()
# check if there are unknown datasets
for dataset in args.datasets.split(','):
if dataset not in datasets_names:
raise ValueError(
'Unsupported or unknown dataset: {}!'.format(dataset))
# check if test dataset are downloaded
# and download if they are not
# download_train(get_data_root())
# download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# loading network from path
if args.network_path is not None:
print(">> Loading network:\n>>>> '{}'".format(args.network_path))
if args.network_path in PRETRAINED:
# pretrained networks (downloaded automatically)
state = load_url(PRETRAINED[args.network_path],
model_dir=os.path.join(get_data_root(),
'networks'))
else:
# fine-tuned network from path
state = torch.load(args.network_path)
# parsing net params from meta
# architecture, pooling, mean, std required
# the rest has default values, in case that is doesnt exist
net_params = {}
net_params['architecture'] = state['meta']['architecture']
net_params['pooling'] = state['meta']['pooling']
net_params['local_whitening'] = state['meta'].get(
'local_whitening', False)
net_params['regional'] = state['meta'].get('regional', False)
net_params['whitening'] = state['meta'].get('whitening', False)
net_params['mean'] = state['meta']['mean']
net_params['std'] = state['meta']['std']
net_params['pretrained'] = False
# load network
net = init_network(net_params)
net.load_state_dict(state['state_dict'])
# if whitening is precomputed
if 'Lw' in state['meta']:
net.meta['Lw'] = state['meta']['Lw']
print(">>>> loaded network: ")
print(net.meta_repr())
# loading offtheshelf network
elif args.network_offtheshelf is not None:
# parse off-the-shelf parameters
offtheshelf = args.network_offtheshelf.split('-')
net_params = {}
net_params['architecture'] = offtheshelf[0]
net_params['pooling'] = offtheshelf[1]
net_params['local_whitening'] = 'lwhiten' in offtheshelf[2:]
net_params['regional'] = 'reg' in offtheshelf[2:]
net_params['whitening'] = 'whiten' in offtheshelf[2:]
net_params['pretrained'] = True
# load off-the-shelf network
print(">> Loading off-the-shelf network:\n>>>> '{}'".format(
args.network_offtheshelf))
net = init_network(net_params)
print(">>>> loaded network: ")
print(net.meta_repr())
# setting up the multi-scale parameters
ms = list(eval(args.multiscale))
if len(ms) > 1 and net.meta['pooling'] == 'gem' and not net.meta[
'regional'] and not net.meta['whitening']:
msp = net.pool.p.item()
print(">> Set-up multiscale:")
print(">>>> ms: {}".format(ms))
print(">>>> msp: {}".format(msp))
else:
msp = 1
# moving network to gpu and eval mode
net.cuda()
net.eval()
# set up the transform
normalize = transforms.Normalize(mean=net.meta['mean'],
std=net.meta['std'])
transform = transforms.Compose([transforms.ToTensor(), normalize])
# evaluate on test datasets
datasets = args.datasets.split(',')
for dataset in datasets:
start = time.time()
print('>> {}: Extracting...'.format(dataset))
print('>> Prepare data information...')
index_file_path = os.path.join(get_data_root(), 'index.csv')
index_mark_path = os.path.join(get_data_root(), 'index_mark.csv')
index_miss_path = os.path.join(get_data_root(), 'index_miss.csv')
test_file_path = os.path.join(get_data_root(), 'test.csv')
test_mark_path = os.path.join(get_data_root(), 'test_mark.csv')
test_mark_add_path = os.path.join(get_data_root(), 'test_mark_add.csv')
test_miss_path = os.path.join(get_data_root(), 'test_miss.csv')
if dataset == 'google-landmarks-dataset':
index_img_path = os.path.join(get_data_root(), 'index')
test_img_path = os.path.join(get_data_root(),
'google-landmarks-dataset-test')
elif dataset == 'google-landmarks-dataset-resize':
index_img_path = os.path.join(get_data_root(),
'resize_index_image')
test_img_path = os.path.join(get_data_root(), 'resize_test_image')
if not (os.path.isfile(index_mark_path)
or os.path.isfile(index_miss_path)):
clear_no_exist(index_file_path, index_mark_path, index_miss_path,
index_img_path)
if not (os.path.isfile(test_mark_path)
or os.path.isfile(test_miss_path)):
clear_no_exist(test_file_path, test_mark_path, test_miss_path,
test_img_path)
print('>> load index image path...')
retrieval_other_dataset = '/home/iap205/Datasets/google-landmarks-dataset-resize'
csvfile = open(index_mark_path, 'r')
csvreader = csv.reader(csvfile)
images = []
miss, add = 0, 0
for line in csvreader:
if line[0] == '1':
images.append(os.path.join(index_img_path, line[1] + '.jpg'))
elif line[0] == '0':
retrieval_img_path = os.path.join(retrieval_other_dataset,
'resize_index_image',
line[1] + '.jpg')
if os.path.isfile(retrieval_img_path):
images.append(retrieval_img_path)
add += 1
miss += 1
csvfile.close()
print(
'>>>> index image miss: {}, supplement: {}, still miss: {}'.format(
miss, add, miss - add))
print('>> load query image path...')
csvfile = open(test_mark_path, 'r')
csvreader = csv.reader(csvfile)
savefile = open(test_mark_add_path, 'w')
save_writer = csv.writer(savefile)
qimages = []
miss, add = 0, 0
for line in csvreader:
if line[0] == '1':
qimages.append(os.path.join(test_img_path, line[1] + '.jpg'))
save_writer.writerow(line)
elif line[0] == '0':
retrieval_img_path = os.path.join(retrieval_other_dataset,
'resize_test_image',
line[1] + '.jpg')
if os.path.isfile(retrieval_img_path):
qimages.append(retrieval_img_path)
save_writer.writerow(['1', line[1]])
add += 1
else:
save_writer.writerow(line)
miss += 1
csvfile.close()
savefile.close()
print(
'>>>> test image miss: {}, supplement: {}, still miss: {}'.format(
miss, add, miss - add))
# extract index vectors
print('>> {}: index images...'.format(dataset))
split_num = 6
extract_num = int(len(images) / split_num)
num_list = list(range(0, len(images), extract_num))
num_list.append(len(images))
# k = 0
# print('>>>> extract part {} of {}'.format(k, split_num-1))
# vecs = extract_vectors(net, images[num_list[k]:num_list[k+1]], args.image_size, transform, ms=ms, msp=msp)
# vecs = vecs.numpy()
# print('>>>> save index vecs to pkl...')
# vecs_file_path = os.path.join(get_data_root(), 'index_vecs{}_of_{}.pkl'.format(k+1, split_num))
# vecs_file = open(vecs_file_path, 'wb')
# pickle.dump(vecs[:, num_list[k]:num_list[k+1]], vecs_file)
# vecs_file.close()
# print('>>>> index_vecs{}_of_{}.pkl save done...'.format(k+1, split_num))
for i in range(split_num):
# vecs_temp = np.loadtxt(open(os.path.join(get_data_root(), 'index_vecs{}_of_{}.csv'.format(i+1, split_num)), "rb"),
# delimiter=",", skiprows=0)
with open(
os.path.join(
get_data_root(),
'index_vecs{}_of_{}.pkl'.format(i + 1, split_num)),
'rb') as f:
vecs_temp = pickle.load(f)
if i == 0:
vecs = vecs_temp
else:
vecs = np.hstack((vecs, vecs_temp[:, :]))
del vecs_temp
gc.collect()
print('\r>>>> index_vecs{}_of_{}.pkl load done...'.format(
i + 1, split_num),
end='')
print('')
# extract query vectors
print('>> {}: query images...'.format(dataset))
split_num = 1
extract_num = int(len(qimages) / split_num)
num_list = list(range(0, len(qimages), extract_num))
num_list.append(len(qimages))
# k = 0
# print('>>>> extract part {} of {}'.format(k, split_num - 1))
# qvecs = extract_vectors(net, qimages[num_list[k]:num_list[k + 1]], args.image_size, transform, ms=ms, msp=msp)
# qvecs = qvecs.numpy()
# for i in range(split_num):
# qvecs_file_path = os.path.join(get_data_root(), 'test_vecs{}_of_{}.pkl'.format(i+1, split_num))
# qvecs_file = open(qvecs_file_path, 'wb')
# pickle.dump(qvecs[:, num_list[i]:num_list[i+1]], qvecs_file)
# qvecs_file.close()
# print('\r>>>> test_vecs{}_of_{}.pkl save done...'.format(i+1, split_num), end='')
# print('')
for i in range(split_num):
# qvecs_temp = np.loadtxt(open(os.path.join(get_data_root(), 'test_vecs{}_of_{}.csv'.format(i+1, split_num)), "rb"),
# delimiter=",", skiprows=0)
with open(
os.path.join(
get_data_root(),
'test_vecs{}_of_{}.pkl'.format(i + 1, split_num)),
'rb') as f:
qvecs_temp = pickle.load(f)
if i == 0:
qvecs = qvecs_temp
else:
qvecs = np.hstack((qvecs, qvecs_temp[:, :]))
del qvecs_temp
gc.collect()
print('\r>>>> test_vecs{}_of_{}.pkl load done...'.format(
i + 1, split_num),
end='')
print('')
# vecs = np.zeros((2048, 1093278))
# qvecs = np.zeros((2048, 115921))
# save vecs to csv file
# np.savetxt(os.path.join(get_data_root(), 'index_vecs{}_of_{}.csv'.format(k, split_num-1)), vecs, delimiter=',')
# np.savetxt(os.path.join(get_data_root(), 'test_vecs{}_of_{}.csv'.format(k, split_num-1)), qvecs, delimiter=',')
# compute whitening
if args.whitening is not None:
start = time.time()
if 'Lw' in net.meta and args.whitening in net.meta['Lw']:
print('>> {}: Whitening is precomputed, loading it...'.format(
args.whitening))
if len(ms) > 1:
Lw = net.meta['Lw'][args.whitening]['ms']
else:
Lw = net.meta['Lw'][args.whitening]['ss']
else:
# if we evaluate networks from path we should save/load whitening
# not to compute it every time
if args.network_path is not None:
whiten_fn = args.network_path + '_{}_whiten'.format(
args.whitening)
if len(ms) > 1:
whiten_fn += '_ms'
whiten_fn += '.pth'
else:
whiten_fn = None
if whiten_fn is not None and os.path.isfile(whiten_fn):
print('>> {}: Whitening is precomputed, loading it...'.
format(args.whitening))
Lw = torch.load(whiten_fn)
else:
print('>> {}: Learning whitening...'.format(
args.whitening))
# extract whitening vectors
print('>> {}: Extracting...'.format(args.whitening))
wvecs = vecs
# learning whitening
print('>> {}: Learning...'.format(args.whitening))
m, P = pcawhitenlearn(wvecs)
Lw = {'m': m, 'P': P}
# saving whitening if whiten_fn exists
if whiten_fn is not None:
print('>> {}: Saving to {}...'.format(
args.whitening, whiten_fn))
torch.save(Lw, whiten_fn)
print('>> {}: elapsed time: {}'.format(args.whitening,
htime(time.time() - start)))
else:
Lw = None
print('>> apply PCAwhiten...')
if Lw is not None:
# whiten the vectors and shrinkage
vecs = whitenapply(vecs, Lw['m'], Lw['P'])
qvecs = whitenapply(qvecs, Lw['m'], Lw['P'])
print('>>>> save index PCAwhiten vecs to pkl...')
split_num = 6
extract_num = int(len(images) / split_num)
num_list = list(range(0, len(images), extract_num))
num_list.append(len(images))
for i in range(split_num):
vecs_file_path = os.path.join(
get_data_root(),
'index_PCAwhiten_vecs{}_of_{}.pkl'.format(i + 1, split_num))
vecs_file = open(vecs_file_path, 'wb')
pickle.dump(vecs[:, num_list[i]:num_list[i + 1]], vecs_file)
vecs_file.close()
print(
'\r>>>> index_PCAwhiten_vecs{}_of_{}.pkl save done...'.format(
i + 1, split_num),
end='')
print('')
print('>>>> save test PCAwhiten vecs to pkl...')
split_num = 1
extract_num = int(len(qimages) / split_num)
num_list = list(range(0, len(qimages), extract_num))
num_list.append(len(images))
for i in range(split_num):
qvecs_file_path = os.path.join(
get_data_root(),
'test_PCAwhiten_vecs{}_of_{}.pkl'.format(i + 1, split_num))
qvecs_file = open(qvecs_file_path, 'wb')
pickle.dump(qvecs[:, num_list[i]:num_list[i + 1]], qvecs_file)
qvecs_file.close()
print('\r>>>> test_PCAwhiten_vecs{}_of_{}.pkl save done...'.format(
i + 1, split_num),
end='')
print('')
print('>>>> load index PCAwhiten vecs from pkl...')
for i in range(split_num):
with open(
os.path.join(
get_data_root(),
'index_PCAwhiten_vecs{}_of_{}.pkl'.format(
i + 1, split_num)), 'rb') as f:
vecs_temp = pickle.load(f)
if i == 0:
vecs = vecs_temp
else:
vecs = np.hstack((vecs, vecs_temp[:, :]))
del vecs_temp
gc.collect()
print(
'\r>>>> index_PCAwhiten_vecs{}_of_{}.pkl load done...'.format(
i + 1, split_num),
end='')
print('')
print('>>>> load test PCAwhiten vecs from pkl...')
for i in range(split_num):
with open(
os.path.join(
get_data_root(),
'test_PCAwhiten_vecs{}_of_{}.pkl'.format(
i + 1, split_num)), 'rb') as f:
qvecs_temp = pickle.load(f)
if i == 0:
qvecs = qvecs_temp
else:
qvecs = np.hstack((qvecs, qvecs_temp[:, :]))
del qvecs_temp
gc.collect()
print('\r>>>> test_PCAwhiten_vecs{}_of_{}.pkl load done...'.format(
i + 1, split_num),
end='')
print('')
# extract principal components and dimension shrinkage
ratio = 0.8
vecs = vecs[:int(vecs.shape[0] * ratio), :]
qvecs = vecs[:int(qvecs.shape[0] * ratio), :]
print('>> {}: Evaluating...'.format(dataset))
split_num = 50
top_num = 100
vecs_T = np.zeros((vecs.shape[1], vecs.shape[0])).astype('float32')
vecs_T[:] = vecs.T[:]
QE_iter = 0
QE_weight = (np.arange(top_num, 0, -1) / top_num).reshape(
1, top_num, 1)
print('>> find {} nearest neighbour...'.format(top_num))
import faiss # place it in the file top will cause network load so slowly
# ranks_top_100 = np.loadtxt(open(os.path.join(get_data_root(), 'ranks_top_{}.csv'.format(top_num)), "rb"),
# delimiter=",", skiprows=0).astype('int')
for iter in range(0, QE_iter + 1):
if iter != 0:
# ranks_top_100 = np.ones((100, 115921)).astype('int')
print('>> Query expansion iteration {}'.format(iter))
ranks_split = 50
for i in range(ranks_split):
ranks_top_100_split = ranks_top_100[:,
int(ranks_top_100.
shape[1] /
ranks_split * i
):int(ranks_top_100
.shape[1] /
ranks_split *
(i + 1))]
top_100_vecs = vecs[:,
ranks_top_100_split] # (2048, 100, query_split_size)
qvecs_temp = (top_100_vecs * QE_weight).sum(axis=1)
qvecs_temp = qvecs_temp / (np.linalg.norm(
qvecs_temp, ord=2, axis=0, keepdims=True) + 1e-6)
if i == 0:
qvecs = qvecs_temp
else:
qvecs = np.hstack((qvecs, qvecs_temp))
del ranks_top_100_split, top_100_vecs, qvecs_temp
gc.collect()
print('\r>>>> calculate new query vectors {}/{} done...'.
format(i + 1, ranks_split),
end='')
print('')
qe_iter_qvecs_path = os.path.join(
get_data_root(), 'QE_iter{}_qvecs.pkl'.format(iter))
qe_iter_qvecs_file = open(qe_iter_qvecs_path, 'wb')
pickle.dump(qvecs, qe_iter_qvecs_file)
qe_iter_qvecs_file.close()
print('>>>> QE_iter{}_qvecs.pkl save done...'.format(iter))
del ranks_top_100
gc.collect()
for i in range(split_num):
# scores = np.dot(vecs.T, qvecs[:, int(qvecs.shape[1]/split_num*i):int(qvecs.shape[1]/split_num*(i+1))])
# ranks = np.argsort(-scores, axis=0)
# kNN search
k = top_num
index = faiss.IndexFlatL2(vecs.shape[0])
index.add(vecs_T)
query_vecs = qvecs[:,
int(qvecs.shape[1] / split_num *
i):int(qvecs.shape[1] / split_num *
(i + 1))]
qvecs_T = np.zeros((query_vecs.shape[1],
query_vecs.shape[0])).astype('float32')
qvecs_T[:] = query_vecs.T[:]
_, ranks = index.search(qvecs_T, k)
ranks = ranks.T
if i == 0:
ranks_top_100 = ranks[:top_num, :]
else:
ranks_top_100 = np.hstack(
(ranks_top_100, ranks[:top_num, :]))
# del scores, ranks
del index, query_vecs, qvecs_T, ranks
gc.collect()
print('\r>>>> kNN search {} nearest neighbour {}/{} done...'.
format(top_num, i + 1, split_num),
end='')
del qvecs
gc.collect()
print('')
del vecs, vecs_T
gc.collect()
# save to csv file
print(">> save to submission.csv file...")
submission_file = open(os.path.join(get_data_root(), 'submission.csv'),
'w')
writer = csv.writer(submission_file)
test_mark_file = open(test_mark_add_path, 'r')
csvreader = csv.reader(test_mark_file)
cnt = 0
writer.writerow(['id', 'images'])
for index, line in enumerate(csvreader):
(flag, img_name) = line[:2]
if flag == '1':
select = []
for i in range(top_num):
select.append(images[int(
ranks_top_100[i,
cnt])].split('/')[-1].split('.jpg')[0])
cnt += 1
writer.writerow([
img_name.split('/')[-1].split('.jpg')[0], ' '.join(select)
])
else:
# random_list = random.sample(range(0, len(images)), top_num)
random_list = np.random.choice(len(images),
top_num,
replace=False)
select = []
for i in range(top_num):
select.append(
images[random_list[i]].split('/')[-1].split('.jpg')[0])
writer.writerow([
img_name.split('/')[-1].split('.jpg')[0], ' '.join(select)
])
if cnt % 10 == 0 or cnt == len(qimages):
print('\r>>>> {}/{} done...'.format(cnt, len(qimages)), end='')
submission_file.close()
test_mark_file.close()
print('')
print('>> {}: elapsed time: {}'.format(dataset,
htime(time.time() - start)))
def main():
global args, min_loss
args = parser.parse_args()
print(args)
# manually check if there are unknown test datasets
for dataset in args.test_datasets.split(','):
if dataset not in test_datasets_names:
raise ValueError(
'Unsupported or unknown test dataset: {}!'.format(dataset))
# check if test dataset are downloaded
# and download if they are not
download_train(get_data_root())
download_test(get_data_root())
# create export dir if it doesnt exist
directory = "{}".format(args.training_dataset)
directory += "_{}".format(args.arch)
directory += "_{}".format(args.pool)
if args.local_whitening:
directory += "_lwhiten"
if args.regional:
directory += "_r"
if args.whitening:
directory += "_whiten"
if not args.pretrained:
directory += "_notpretrained"
directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin)
directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr,
args.weight_decay)
directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num,
args.query_size,
args.pool_size)
directory += "_bsize{}_uevery{}_imsize{}".format(args.batch_size,
args.update_every,
args.image_size)
args.directory = os.path.join(args.directory, directory)
print(">> Creating directory if it does not exist:\n>> '{}'".format(
args.directory))
if not os.path.exists(args.directory):
os.makedirs(args.directory)
# set cuda visible device
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# set random seeds
# TODO: maybe pass as argument in future implementation?
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
np.random.seed(0)
# initialize model
if args.pretrained:
print(">> Using pre-trained model '{}'".format(args.arch))
else:
print(">> Using model from scratch (random weights) '{}'".format(
args.arch))
model_params = {}
model_params['architecture'] = args.arch
model_params['pooling'] = args.pool
model_params['local_whitening'] = args.local_whitening
model_params['regional'] = args.regional
model_params['whitening'] = args.whitening
model_params['mean'] = args.mean
model_params['std'] = args.std
model_params['pretrained'] = args.pretrained
model = init_network(model_params)
# move network to gpu
model.cuda()
print(model)
summary(model, (3, 224, 224))
# define loss function (criterion) and optimizer
if args.loss == 'contrastive':
criterion = ContrastiveLoss(margin=args.loss_margin).cuda()
elif args.loss == 'triplet':
criterion = TripletLoss(margin=args.loss_margin).cuda()
else:
raise (RuntimeError("Loss {} not available!".format(args.loss)))
# parameters split into features, pool, whitening
# IMPORTANT: no weight decay for pooling parameter p in GeM or regional-GeM
parameters = []
# add feature parameters
parameters.append({'params': model.features.parameters()})
# add local whitening if exists
if model.lwhiten is not None:
parameters.append({'params': model.lwhiten.parameters()})
# add pooling parameters (or regional whitening which is part of the pooling layer!)
if not args.regional:
# global, only pooling parameter p weight decay should be 0
if args.pool == 'gem':
parameters.append({
'params': model.pool.parameters(),
'lr': args.lr * 10,
'weight_decay': 0
})
elif args.pool == 'gemmp':
parameters.append({
'params': model.pool.parameters(),
'lr': args.lr * 100,
'weight_decay': 0
})
else:
# regional, pooling parameter p weight decay should be 0,
# and we want to add regional whitening if it is there
if args.pool == 'gem':
parameters.append({
'params': model.pool.rpool.parameters(),
'lr': args.lr * 10,
'weight_decay': 0
})
elif args.pool == 'gemmp':
parameters.append({
'params': model.pool.rpool.parameters(),
'lr': args.lr * 100,
'weight_decay': 0
})
if model.pool.whiten is not None:
parameters.append({'params': model.pool.whiten.parameters()})
# add final whitening if exists
if model.whiten is not None:
parameters.append({'params': model.whiten.parameters()})
# define optimizer
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(parameters,
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(parameters,
args.lr,
weight_decay=args.weight_decay)
# define learning rate decay schedule
# TODO: maybe pass as argument in future implementation?
exp_decay = math.exp(-0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer,
gamma=exp_decay)
# optionally resume from a checkpoint
start_epoch = 0
if args.resume:
args.resume = os.path.join(args.directory, args.resume)
if os.path.isfile(args.resume):
# load checkpoint weights and update model and optimizer
print(">> Loading checkpoint:\n>> '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
min_loss = checkpoint['min_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# important not to forget scheduler updating
scheduler = torch.optim.lr_scheduler.ExponentialLR(
optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch'] - 1)
else:
print(">> No checkpoint found at '{}'".format(args.resume))
# Data loading code
normalize = transforms.Normalize(mean=model.meta['mean'],
std=model.meta['std'])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
train_dataset = TuplesDataset(name=args.training_dataset,
mode='train',
imsize=args.image_size,
nnum=args.neg_num,
qsize=args.query_size,
poolsize=args.pool_size,
transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
sampler=None,
drop_last=True,
collate_fn=collate_tuples)
if args.val:
val_dataset = TuplesDataset(name=args.training_dataset,
mode='val',
imsize=args.image_size,
nnum=args.neg_num,
qsize=float('Inf'),
poolsize=float('Inf'),
transform=transform)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
drop_last=True,
collate_fn=collate_tuples)
# uncomment to evaluate the network before starting
# test(args.test_datasets, model)
for epoch in range(start_epoch, args.epochs):
# set manual seeds per epoch
np.random.seed(epoch)
torch.manual_seed(epoch)
torch.cuda.manual_seed_all(epoch)
# adjust learning rate for each epoch
scheduler.step()
# # debug printing to check if everything ok
# lr_feat = optimizer.param_groups[0]['lr']
# lr_pool = optimizer.param_groups[1]['lr']
# print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool))
# train for one epoch on train set
# loss = train(train_loader, model, criterion, optimizer, epoch)
loss = 999.
# evaluate on validation set
if args.val:
with torch.no_grad():
loss = validate(val_loader, model, criterion, epoch)
# evaluate on test datasets every test_freq epochs
if (epoch + 1) % args.test_freq == 0:
with torch.no_grad():
test(args.test_datasets, model)
# remember best loss and save checkpoint
is_best = loss < min_loss
min_loss = min(loss, min_loss)
save_checkpoint(
{
'epoch': epoch + 1,
'meta': model.meta,
'state_dict': model.state_dict(),
'min_loss': min_loss,
'optimizer': optimizer.state_dict(),
}, is_best, args.directory)
def main():
global args, min_loss
args = parser.parse_args()
# manually check if there are unknown test datasets
for dataset in args.test_datasets.split(','):
if dataset not in test_datasets_names:
raise ValueError('Unsupported or unknown test dataset: {}!'.format(dataset))
# check if test dataset are downloaded
# and download if they are not
# download_train(get_data_root())
# download_test(get_data_root())
# download_train('/media/hxq/ExtractedDatasets/Datasets/retrieval-SfM')
# download_test('/media/hxq/ExtractedDatasets/Datasets/retrieval-SfM')
# create export dir if it doesnt exist
directory = "{}".format(args.training_dataset)
directory += "_{}".format(args.arch)
directory += "_{}".format(args.pool)
if args.local_whitening:
directory += "_lwhiten"
if args.regional:
directory += "_r"
if args.whitening:
directory += "_whiten"
if not args.pretrained:
directory += "_notpretrained"
directory += "_{}_m{:.2f}".format(args.loss, args.loss_margin)
directory += "_{}_lr{:.1e}_wd{:.1e}".format(args.optimizer, args.lr, args.weight_decay)
directory += "_nnum{}_qsize{}_psize{}".format(args.neg_num, args.query_size, args.pool_size)
directory += "_bsize{}_imsize{}".format(args.batch_size, args.image_size)
args.directory = os.path.join(args.directory, directory)
print(">> Creating directory if it does not exist:\n>> '{}'".format(args.directory))
if not os.path.exists(args.directory):
os.makedirs(args.directory)
# set cuda visible device
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# set random seeds
# TODO: maybe pass as argument in future implementation?
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
np.random.seed(0)
# initialize model
if args.pretrained:
print(">> Using pre-trained model '{}'".format(args.arch))
else:
print(">> Using model from scratch (random weights) '{}'".format(args.arch))
model_params = {}
model_params['architecture'] = args.arch
model_params['pooling'] = args.pool
model_params['local_whitening'] = args.local_whitening
model_params['regional'] = args.regional
model_params['whitening'] = args.whitening
# model_params['mean'] = ... # will use default
# model_params['std'] = ... # will use default
model_params['pretrained'] = args.pretrained
model_params['multi_layer_cat'] = args.multi_layer_cat
model = init_network(model_params)
print(">>>> loaded model: ")
print(model.meta_repr())
# move network to gpu
model.cuda()
# define loss function (criterion) and optimizer
if args.loss == 'contrastive':
criterion = ContrastiveLoss(margin=args.loss_margin).cuda()
else:
raise(RuntimeError("Loss {} not available!".format(args.loss)))
# parameters split into features, pool, whitening
# IMPORTANT: no weight decay for pooling parameters p in GeM or regional-GeM
parameters = []
# add feature parameters
parameters.append({'params': model.features.parameters()})
# add local whitening if exists
if model.lwhiten is not None:
parameters.append({'params': model.lwhiten.parameters()})
# add pooling parameters (or regional whitening which is part of the pooling layer!)
if not args.regional:
# global, only pooling parameters p weight decay should be 0
parameters.append({'params': model.pool.parameters(), 'lr': args.lr*10, 'weight_decay': 0})
else:
# regional, pooling parameters p weight decay should be 0,
# and we want to add regional whitening if it is there
parameters.append({'params': model.pool.rpool.parameters(), 'lr': args.lr*10, 'weight_decay': 0})
if model.pool.whiten is not None:
parameters.append({'params': model.pool.whiten.parameters()})
# add final whitening if exists
if model.whiten is not None:
parameters.append({'params': model.whiten.parameters()})
# define optimizer
if args.optimizer == 'sgd':
optimizer = torch.optim.SGD(parameters, args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(parameters, args.lr, weight_decay=args.weight_decay)
# define learning rate decay schedule
# TODO: maybe pass as argument in future implementation?
exp_decay = math.exp(-0.01)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay)
# optionally resume from a checkpoint
start_epoch = 0
if args.resume:
args.resume = os.path.join(args.directory, args.resume)
if os.path.isfile(args.resume):
# load checkpoint weights and update model and optimizer
print(">> Loading checkpoint:\n>> '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
# hxq added, for whitening test
# print('>> load the parameters of supervised whitening for the FC layer initialization')
# whiten_fn = '/media/iap205/Data/Export/cnnimageretrieval-google_landmark_retrieval/trained_network/' \
# 'R101_O_GL_FC/google-landmarks-dataset-resize_resnet101_gem_whiten_contrastive_m0.85_' \
# 'adam_lr5.0e-07_wd1.0e-04_nnum5_qsize2000_psize22000_bsize5_imsize362/' \
# 'model_epoch114.pth.tar_google-landmarks-dataset_whiten_ms.pth'
# Lw = torch.load(whiten_fn)
# P = Lw['P']
# m = Lw['m']
# P = torch.from_numpy(P).float()
# m = torch.from_numpy(m).float()
# checkpoint['state_dict']['whiten.weight'] = P
# checkpoint['state_dict']['whiten.bias'] = -torch.mm(P, m).squeeze()
min_loss = checkpoint['min_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print(">>>> loaded checkpoint:\n>>>> '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
# important not to forget scheduler updating
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=exp_decay, last_epoch=checkpoint['epoch']-1)
# print(">> Use new scheduler")
else:
print(">> No checkpoint found at '{}'".format(args.resume))
# Data loading
# hxq added, but it can be generate in advance
# if args.training_dataset == 'google-landmarks-dataset-v2':
# data_table_path = os.path.join(get_data_root(), 'train.csv')
# img_check_path = os.path.join(get_data_root(), 'train')
# train_val_file_output_path = os.path.join(get_data_root())
# test_file_output_path = os.path.join(get_data_root(), 'test', args.training_dataset+'-test')
# val_set_size = 10000
# test_set_size = 10000
# elif args.training_dataset == 'google-landmarks-dataset-resize':
# data_table_path = os.path.join(get_data_root(), 'train.csv')
# img_check_path = os.path.join(get_data_root(), 'resize_train_image')
# train_val_file_output_path = os.path.join(get_data_root())
# test_file_output_path = os.path.join(get_data_root(), 'test', args.training_dataset+'-test')
# val_set_size = 10000
# test_set_size = 10000
# elif args.training_dataset == 'google-landmarks-dataset':
# pass
# if args.training_dataset == 'google-landmarks-dataset-v2'\
# or args.training_dataset == 'google-landmarks-dataset-resize'\
# or args.training_dataset == 'google-landmarks-dataset':
# if not (os.path.isfile(os.path.join(train_val_file_output_path, '{}.pkl'.format(args.training_dataset)))
# and os.path.isfile(os.path.join(test_file_output_path, 'gnd_{}-test.pkl'.format(args.training_dataset)))):
# gen_train_val_test_pkl(args.training_dataset, data_table_path, img_check_path, val_set_size, test_set_size,
# train_val_file_output_path, test_file_output_path)
normalize = transforms.Normalize(mean=model.meta['mean'], std=model.meta['std'])
transform = transforms.Compose([
transforms.ToTensor(),
normalize,
])
train_dataset = TuplesDataset(
name=args.training_dataset,
mode='train',
imsize=args.image_size,
nnum=args.neg_num,
qsize=args.query_size,
poolsize=args.pool_size,
transform=transform
)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True, sampler=None,
drop_last=True, collate_fn=collate_tuples
)
if args.val:
val_dataset = TuplesDataset(
name=args.training_dataset,
mode='val',
imsize=args.image_size,
nnum=args.neg_num,
qsize=float('Inf'),
poolsize=float('Inf'),
transform=transform
)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True,
drop_last=True, collate_fn=collate_tuples
)
# evaluate the network before starting
# this might not be necessary?
# test(args.test_datasets, model)
# add by hxq
# save_checkpoint({
# 'epoch': 0,
# 'meta': model.meta,
# 'state_dict': model.state_dict(),
# 'min_loss': min_loss,
# 'optimizer': optimizer.state_dict(),
# }, 0, args.directory)
# hxq added, save trained network as epoch
# network_path = {'vgg16': '/home/iap205/Datasets/retrieval-SfM/networks/retrievalSfM120k-vgg16-gem-b4dcdc6.pth',
# 'resnet101': '/home/iap205/Datasets/retrieval-SfM/networks/retrievalSfM120k-resnet101-gem-b80fb85.pth'}
# model_trained = torch.load(network_path['resnet101'])
# save_checkpoint({
# 'epoch': 0,
# 'meta': model.meta,
# 'state_dict': model_trained['state_dict'],
# 'min_loss': min_loss,
# 'optimizer': optimizer.state_dict(),
# }, 0, args.directory)
for epoch in range(start_epoch, args.epochs):
# set manual seeds per epoch
np.random.seed(epoch)
torch.manual_seed(epoch)
torch.cuda.manual_seed_all(epoch)
# adjust learning rate for each epoch
scheduler.step()
# # debug printing to check if everything ok
# lr_feat = optimizer.param_groups[0]['lr']
# lr_pool = optimizer.param_groups[1]['lr']
# print('>> Features lr: {:.2e}; Pooling lr: {:.2e}'.format(lr_feat, lr_pool))
# train for one epoch on train set
loss = train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
if args.val:
with torch.no_grad():
loss = validate(val_loader, model, criterion, epoch)
# evaluate on test datasets every test_freq epochs
if (epoch + 1) % args.test_freq == 0:
with torch.no_grad():
test(args.test_datasets, model)
# remember best loss and save checkpoint
is_best = loss < min_loss
min_loss = min(loss, min_loss)
save_checkpoint({
'epoch': epoch + 1,
'meta': model.meta,
'state_dict': model.state_dict(),
'min_loss': min_loss,
'optimizer' : optimizer.state_dict(),
}, is_best, args.directory)
