def get_style_model_and_losses(cnn,
style_img,
content_img,
style_weight=1000,
content_weight=1,
content_layers=content_layers_default,
style_layers=style_layers_default):
cnn = copy.deepcopy(cnn)
# just in order to have an iterable access to or list of content/syle
# losses
content_losses = []
style_losses = []
model = nn.Sequential() # the new Sequential module network
gram = GramMatrix(
) # we need a gram module in order to compute style targets
# move these modules to the GPU if possible:
if use_cuda:
model = model.cuda()
gram = gram.cuda()
i = 1
for layer in list(cnn):
if isinstance(layer, nn.Conv2d):
name = "conv_" + str(i)
model.add_module(name, layer)
if name in content_layers:
# add content loss:
target = model(content_img).clone()
content_loss = ContentLoss(target, content_weight)
model.add_module("content_loss_" + str(i), content_loss)
content_losses.append(content_loss)
if name in style_layers:
# add style loss:
target_feature = model(style_img).clone()
target_feature_gram = gram(target_feature)
style_loss = StyleLoss(target_feature_gram, style_weight)
model.add_module("style_loss_" + str(i), style_loss)
style_losses.append(style_loss)
if isinstance(layer, nn.ReLU):
name = "relu_" + str(i)
model.add_module(name, layer)
if name in content_layers:
# add content loss:
target = model(content_img).clone()
content_loss = ContentLoss(target, content_weight)
model.add_module("content_loss_" + str(i), content_loss)
content_losses.append(content_loss)
if name in style_layers:
# add style loss:
target_feature = model(style_img).clone()
target_feature_gram = gram(target_feature)
style_loss = StyleLoss(target_feature_gram, style_weight)
model.add_module("style_loss_" + str(i), style_loss)
style_losses.append(style_loss)
i += 1
if isinstance(layer, nn.MaxPool2d):
name = "pool_" + str(i)
model.add_module(name, layer) # ***
return model, style_losses, content_losses
ip_tfs = tf.Compose([tf.ToTensor()])
pil_tf = tf.ToPILImage()
vgg19 = models.vgg19(pretrained=True).features
model = nn.Sequential()
gm = GramMatrix()
dtype = torch.FloatTensor
if torch.cuda.is_available():
dtype = torch.cuda.FloatTensor
if args.cuda:
vgg19 = vgg19.cuda()
model = model.cuda()
gm = gm.cuda()
content_layers = [
'conv_4'
] # The fourth convolutional layer's activation to serve as the content
style_layers = [
'conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5'
] # The list of layers contributing to the style representations
cl_list = [] # to compute the common loss term
sl_list = []
photo = imgloader(args.photo_dir, ip_tfs, 512).type(dtype)
art = imgloader(args.art_dir, ip_tfs, 512).type(dtype)
print(photo.size())
print(art.size())