def define_P(gpu_ids=[]):
use_gpu = len(gpu_ids) > 0
if use_gpu:
assert (torch.cuda.is_available())
netP = Vgg16()
util.init_vgg16('./')
netP.load_state_dict(torch.load(os.path.join('./', "vgg16.weight")))
for param in netP.parameters():
param.requires_grad = False
if use_gpu:
netP.cuda()
return netP
def define_P(perceptual_model_dir, gpu_ids=[]):
use_gpu = len(gpu_ids) > 0
if use_gpu:
assert (torch.cuda.is_available())
netP = Vgg16()
util.init_vgg16(perceptual_model_dir)
netP.load_state_dict(
torch.load(os.path.join(perceptual_model_dir, "vgg16.weight")))
for param in netP.parameters():
param.requires_grad = False
if use_gpu:
netP.cuda(device=gpu_ids[0])
return netP
opt.nThreads = 1 # test code only supports nThreads=1
opt.batchSize = 1 #test code only supports batchSize=1
opt.serial_batches = True # no shuffle
opt.lambda_p = 0
# Load vocabulary wrapper.
with open(opt.vocab_path, 'rb') as f:
vocab = pickle.load(f)
opt.vocab = vocab
opt.vocab_size = len(vocab)
print('load vgg16 models')
init_vgg16("vgg_model")
vgg_model = Vgg16Part()
vgg_model.load_state_dict(torch.load('vgg_model/vgg16.weight'))
if torch.cuda.is_available():
vgg_model.cuda()
opt.vgg_model = vgg_model
data_loader = CreateDataLoader(opt)
dataset = data_loader.load_data()
model = create_model(opt)
visualizer = Visualizer(opt)
# create website
web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))
# test