def get_model(opt):
base = base_network.VGG(pretrained=True)
net = SiameseNetwork(base=base)
if opt.mode == 'train' and not opt.continue_train:
net.fc1.apply(weights_init)
net.fc2.apply(weights_init)
net.fc3.apply(weights_init)
net.fc_linear.weight.data.fill_(1.0)
else:
# HACK: strict=False
net.load_state_dict(torch.load(
os.path.join(opt.checkpoint_dir, opt.name,
'{}_net.pth'.format(opt.which_epoch))),
strict=True)
if opt.mode != 'train':
set_requires_grad(net, False)
net.eval()
if opt.use_gpu:
net.cuda()
return net
def __init__(self, pretrained=True):
super(encoder, self).__init__()
self.model = base_network.VGG(pretrained=pretrained)
def _get_aux_nets(self):
self.vgg_teacher = nn.DataParallel(
base_network.VGG(pretrained=True)).cuda()
self.perceptural_loss = perceptural_loss().cuda()
self.KGTransform = nn.DataParallel(nn.Conv2d(512, 512, 1)).cuda()
def __init__(self):
super(perceptural_loss, self).__init__()
self.vgg = base_network.VGG(pretrained=True).eval().cuda()
self.vgg.features_1 = nn.DataParallel(self.vgg.features_1)
self.mse = nn.MSELoss()