def get_model(args):
base_model = models.create(args.arch, train_layers=args.layers, matconvnet='logs/vd16_offtheshelf_conv5_3_max.pth')
if args.vlad:
pool_layer = models.create('netvlad', dim=base_model.feature_dim)
# vgg16_pitts_64_desc_cen_mat.hdf5
initcache = osp.join(args.init_dir, args.arch + '_' + args.dataset + '_' + str(args.num_clusters) + '_desc_cen.hdf5')
if (dist.get_rank()==0):
print ('Loading centroids from {}'.format(initcache))
with h5py.File(initcache, mode='r') as h5:
pool_layer.clsts = h5.get("centroids")[...]
pool_layer.traindescs = h5.get("descriptors")[...]
pool_layer._init_params()
model = models.create('embednet', base_model, pool_layer)
else:
model = base_model
if (args.syncbn):
# not work for VGG16
convert_sync_bn(model)
model.cuda(args.gpu)
model = nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu], output_device=args.gpu, find_unused_parameters=True
)
return model
def vgg16_netvlad(pretrained=False):
base_model = models.create('vgg16', pretrained=False)
pool_layer = models.create('netvlad', dim=base_model.feature_dim)
model = models.create('embednetpca', base_model, pool_layer)
if pretrained:
model.load_state_dict(
torch.hub.load_state_dict_from_url(
'https://github.com/yxgeee/OpenIBL/releases/download/v0.1.0-beta/vgg16_netvlad.pth',
map_location=torch.device('cpu')))
return model
def get_model(args):
base_model = models.create(args.arch)
if args.vlad:
pool_layer = models.create('netvlad', dim=base_model.feature_dim)
model = models.create('embednet', base_model, pool_layer)
else:
model = base_model
model.cuda(args.gpu)
model = nn.parallel.DistributedDataParallel(model,
device_ids=[args.gpu],
output_device=args.gpu,
find_unused_parameters=True)
return model
def get_model(args):
model = models.create(args.arch,
pretrained=True,
cut_at_pooling=True,
matconvnet='logs/vd16_offtheshelf_conv5_3_max.pth')
model.cuda()
model = nn.DataParallel(model)
return model