def download_datasets(datasets, dataset_url, globals):
"""Download data associated with each required dataset"""
if "val_eccv20" in datasets:
download_train(globals['root_path'])
io_helpers.download_files(["retrieval-SfM-120k-val-eccv2020.pkl"],
globals['root_path'] / "train/retrieval-SfM-120k",
dataset_url, logfunc=globals["logger"].info)
elif "train" in datasets:
download_train(globals['root_path'])
if "roxford5k" in datasets or "rparis6k" in datasets:
download_test(globals['root_path'])
def configdataset():
#dataset = dataset.lower()
# if dataset not in DATASETS:
# raise ValueError('Unknown dataset: {}!'.format(dataset))
# loading imlist, qimlist, and gnd, in cfg as a dict
#dataset's path
#gnd_fname = os.path.join(dir_main, dataset, 'gnd_{}.pkl'.format(dataset))
# gnd_fname = os.path.join(dir_main, dataset, 'gnd_{}.pkl'.format(dataset))
# with open(gnd_fname, 'rb') as f:
# cfg = pickle.load(f)
# cfg['gnd_fname'] = gnd_fname
cfg = download_test()
# cfg['ext'] = '.jpg'
# cfg['qext'] = '.jpg'
#cfg['dir_data'] = os.path.join(dir_main, dataset)
cfg['dir_images'] = cfg['png_file']
cfg['gnd'] = cfg['l_file']
cfg['n'] = len(cfg['imlist'])
cfg['nq'] = len(cfg['qimlist'])
#file path
cfg['im_fname'] = config_imname
cfg['qim_fname'] = config_qimname
return cfg
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)))
import os from cirtorch.utils.download import download_test from cirtorch.utils.general import get_root data_root = os.path.join(get_root(), 'cirtorch') download_test(data_root)
def make_dataloader(args, config, rank=None, world_size=None):
general_config = config["general"]
data_config = config["dataloader"]
# Manually check if there are unknown test datasets
for dataset in data_config.getstruct("test_datasets"):
if dataset not in test_datasets_names:
raise ValueError(
'Unsupported or unknown test dataset: {}!'.format(dataset))
# Check if test datasets are available, download it if not !
download_train(args.data)
download_test(args.data)
# Data Loader
log_debug("Creating dataloaders for dataset in %s", args.data)
train_tf = TuplesTransform(
shortest_size=data_config.getint("train_shortest_size"),
longest_max_size=data_config.getint("train_longest_max_size"),
rgb_mean=data_config.getstruct("rgb_mean"),
rgb_std=data_config.getstruct("rgb_std"),
random_flip=data_config.getboolean("random_flip"),
random_scale=data_config.getstruct("random_scale"))
# Training dataloader
train_db = TuplesDataset(root_dir=args.data,
name=data_config.get("training_dataset"),
mode='train',
neg_num=data_config.getint("neg_num"),
query_size=data_config.getint("train_query_size"),
pool_size=data_config.getint("train_pool_size"),
transform=train_tf,
batch_size=data_config.getint("train_batch_size"),
num_workers=data_config.getint("num_workers"))
train_sampler = TuplesDistributedARBatchSampler(
train_db,
batch_size=data_config.getint("train_batch_size"),
num_replicas=world_size,
rank=rank,
drop_last=True)
train_dl = data.DataLoader(train_db,
batch_sampler=train_sampler,
collate_fn=iss_collate_fn,
pin_memory=True,
num_workers=data_config.getint("num_workers"),
shuffle=False)
# Validation dataloader
val_tf = TuplesTransform(
shortest_size=data_config.getint("train_shortest_size"),
longest_max_size=data_config.getint("train_longest_max_size"),
rgb_mean=data_config.getstruct("rgb_mean"),
rgb_std=data_config.getstruct("rgb_std"),
random_flip=data_config.getboolean("random_flip"),
random_scale=data_config.getstruct("random_scale"))
val_db = TuplesDataset(
root_dir=args.data,
name=data_config.get("training_dataset"),
mode='val',
neg_num=data_config.getint("neg_num"),
query_size=float('inf'),
pool_size=float('inf'),
#query_size=data_config.getint("train_query_size"),
#pool_size=data_config.getint("train_pool_size"),
transform=val_tf,
batch_size=data_config.getint("train_batch_size"),
num_workers=data_config.getint("num_workers"))
val_sampler = TuplesDistributedARBatchSampler(
val_db,
batch_size=data_config.getint("train_batch_size"),
num_replicas=world_size,
rank=rank,
drop_last=True)
val_dl = data.DataLoader(val_db,
batch_sampler=val_sampler,
collate_fn=iss_collate_fn,
pin_memory=True,
num_workers=data_config.getint("num_workers"),
shuffle=False)
return train_dl, val_dl
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)
import os.path
import sys
import yaml
sys.path.append(os.path.abspath(__file__ + "/../../../../"))
import mdir
from daan.core.experiments import dict_deep_overlay
# Download necessary datasets
from cirtorch.utils.download import download_test
from cirtorch.utils.general import get_data_root
download_test(get_data_root())
# Parse arguments
scenarios = sys.argv[1:]
if len(scenarios) == 1 and not scenarios[0].endswith(".yml"):
scenarios = ["eval.yml", "eval_%s.yml" % scenarios[0]]
# Parse input scenarios
scenario = {}
for params in scenarios:
with open(params, 'r') as handle:
scenario = dict_deep_overlay(scenario, yaml.safe_load(handle))
if not scenario:
sys.stderr.write("Scenario needs to be specified\n")
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():
print(">> Retrieval evaluation of attacks\n")
args = parser.parse_args()
# check if unknown dataset
if args.dataset not in datasets_names:
raise ValueError('Unsupported or unknown dataset: {}!'.format(
args.dataset))
# check if test dataset are downloaded and download if they are not
download_test(get_data_root())
# setting up the visible GPU
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id
# parse off-the-shelf parameters
offtheshelf = args.network_offtheshelf.split('-')
net_params = {
'architecture': offtheshelf[0],
'pooling': offtheshelf[1],
'local_whitening': False,
'regional': False,
'whitening': False,
'pretrained': True
}
# load off-the-shelf network
print(">> Loading off-the-shelf network: '{}'".format(
args.network_offtheshelf))
net = init_network(net_params)
# print(">>>> loaded network: \n'{}'".format(net.meta_repr()))
# 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 dataset
dataset = args.dataset
print('>> {}: Extracting...'.format(dataset))
# prepare config structure for the test dataset
cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
cfg.update({'qext_a': args.ext_attack})
cfg.update({'dir_data_a': args.dir_attack})
cfg.update({'dir_images_a': cfg['dir_data_a']})
cfg.update({'qim_fname_a': config_qimname_a})
# reduce number of queries for holidays and copydays
if dataset.startswith('holidays') or dataset.startswith('copydays'):
cfg['nq'] = 50
cfg['gnd'] = cfg['gnd'][:cfg['nq']]
images = [cfg['im_fname'](cfg, i) for i in range(cfg['n'])]
qimages = [cfg['qim_fname'](cfg, i) for i in range(cfg['nq'])]
qimages_a = [cfg['qim_fname_a'](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 holidays and copydays
# extract descriptors and cache or load cached ones
print('>> {}: database images...'.format(dataset))
network_fn = args.network_offtheshelf
if args.dir_cache is not None:
vecs_fn = os.path.join(
args.dir_cache,
'{}_{}_{}_{}_vecs.pth'.format(dataset, network_fn, args.image_size,
args.image_resize))
if os.path.isfile(vecs_fn):
vecs = torch.load(vecs_fn)
print('>> loaded cached descriptors from {}'.format(vecs_fn))
else:
vecs = extract_vectors_a(net, images, args.image_size,
args.image_resize, transform)
torch.save(vecs, vecs_fn)
print('>> {}: standard query images...'.format(dataset))
qvecs = extract_vectors_a(net,
qimages,
args.image_size,
args.image_resize,
transform,
bbxs=bbxs)
print('>> {}: attack query images...'.format(dataset))
qvecs_a = extract_vectors_a(net, qimages_a, args.image_size,
args.image_resize, transform)
print('>> {}: evaluating for image resolution {}'.format(
dataset, args.image_resize))
# convert to numpy
vecs = vecs.numpy()
qvecs = qvecs.numpy()
qvecs_a = qvecs_a.numpy()
qip = (qvecs * qvecs_a).sum(axis=0)
qip_mean, qip_std = qip.mean(), qip.std()
print('>> {}: inner product (target,attack) mean: {:.3f}, std: {:.3f}'.
format(dataset, qip_mean, qip_std))
# search, rank, and print
scores = np.dot(vecs.T, qvecs)
ranks = np.argsort(-scores, axis=0)
maps, mprs, aps = compute_map_and_print(dataset, ranks, cfg['gnd'])
# attack search, rank, and print
scores = np.dot(vecs.T, qvecs_a)
ranks = np.argsort(-scores, axis=0)
maps_a, mprs_a, aps_a = compute_map_and_print(
dataset + ' attack ({})'.format(args.ext_attack), ranks, cfg['gnd'])
if dataset.startswith('roxford5k') or dataset.startswith('rparis6k'):
r1, r2, r3 = 100 * maps[1], 100 * maps_a[1], 100 * (
maps_a[1] - maps[1]) # medium protocol
else:
r1, r2, r3 = 100 * maps[0], 100 * maps_a[0], 100 * (maps_a[0] -
maps[0])
print(
'\n*** Summary ***\n attack: {}\n test: {}-{}-{} \n mean ip (target,attack): {:.3f}\n mAP: org {:.2f} att {:.2f} dif {:.2f}\n'
.format(
args.dir_attack.split('/')[-2], dataset, args.network_offtheshelf,
args.image_resize, qip_mean, r1, r2, r3))
if args.log is not None:
with open(args.log, 'a') as f:
f.write(
'\n attack: {}\n test: {}-{}-{} \n mean ip (target,attack): {:.3f}\n mAP: org {:.2f} att {:.2f} dif {:.2f}\n\n'
.format(
args.dir_attack.split('/')[-2], dataset,
args.network_offtheshelf, args.image_resize, qip_mean, r1,
r2, r3))
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():
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():
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:
net = load_network(args.network_path)
# loading offtheshelf network
elif args.network_offtheshelf is not None:
net = load_offtheshelf(args.network_offtheshelf)
# 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))
if args.whitening == "scores":
# special logic for scores database
from score_retrieval.exports import (
db,
train_images as images,
)
else:
# loading db
db_root = os.path.join(get_data_root(), 'train', args.test_whiten)
ims_root = os.path.join(db_root, 'ims')
db_fn = os.path.join(db_root, '{}-whiten.pkl'.format(args.test_whiten))
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))
if dataset == "scores":
# Special added logic to handle loading our score dataset
from score_retrieval.exports import (
images,
qimages,
gnd,
)
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, ms=ms, msp=msp)
else:
# extract ground truth
cfg = configdataset(dataset, os.path.join(get_data_root(), 'test'))
gnd = cfg['gnd']
# prepare config structure for the test dataset
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(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)
# validation
print(">> {}: gnd stats: {}, {}, {}".format(
dataset,
len(gnd),
[len(x["ok"]) for x in gnd[10:]],
[len(x["junk"]) for x in gnd[10:]],
))
print(">> {}: image stats: {}, {}".format(dataset, len(images), len(qimages)))
assert len(gnd) == len(qimages), (len(gnd), len(qimages))
print('>> {}: Evaluating...'.format(dataset))
# convert to numpy
vecs = vecs.numpy()
qvecs = qvecs.numpy()
print(">> {}: qvecs.shape: {}".format(dataset, qvecs.shape))
# search, rank, and print
scores = np.dot(vecs.T, qvecs)
ranks = np.argsort(-scores, axis=0)
print(">> {}: ranks (shape {}) head: {}".format(dataset, ranks.shape, ranks[10:,10:]))
print(">> {}: gnd head: {}".format(dataset, gnd[5:]))
# Compute and print metrics
compute_acc(ranks, gnd, dataset)
compute_mrr(ranks, gnd, dataset)
compute_map_and_print(dataset, ranks, 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_acc(ranks, gnd, dataset + " + whiten")
compute_mrr(ranks, gnd, dataset + " + whiten")
compute_map_and_print(dataset + " + whiten", ranks, 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])
# 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():
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():
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)
