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
help='Content weight')
parser.add_argument('--w_style', type=float, default=1e3, help='Style weight')
args = parser.parse_args()
check_dirs(args.photo_dir)
check_dirs(args.art_dir)
check_dirs(args.out_dir, make=True)
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