ImageSetup.py

import torch
from torch.autograd import Variable
from PIL import Image
import torchvision
import torchvision.transforms as transforms
import torchvision.models as models
from torchvision.utils import save_image

def ImageLoader(image_name, image_size):
    image = Image.open(image_name)
    loader = transforms.Compose([transforms.Scale(image_size), transforms.ToTensor()])  # resize and convert to tensor
    #image = loader(image)
    image = Variable(loader(image))
    # fake batch dimension required to fit network's input dimensions
    image = image.unsqueeze(0)
    return image

def SaveImage(output_img, output_name): 
   torchvision.utils.save_image(output_img, output_name, nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0)

LossModules.py

import torch
import torch.legacy.nn as nn
import torchvision

# Define an nn Module to compute content loss in-place
class ContentLoss(nn.Module):
    def __init__(self, strength, normalize):
      super(ContentLoss, self).__init__()
      self.strength = strength
      self.target = torch.Tensor()
      self.normalize = 'false'
      self.loss = 0
      self.crit = nn.MSECriterion()
      self.mode = None
          
    def updateOutput(self, input):
      if self.mode == 'loss':
        self.loss = self.crit.updateOutput(input, self.target) * self.strength #Forward
      elif self.mode == 'capture':
        self.target.resize_as_(input).copy_(input)
      self.output = input 
      return self.output
  
    def updateGradInput(self, input, gradOutput):
      if self.mode == 'loss':
        if input.nelement() == self.target.nelement():
          self.gradInput = self.crit.updateGradInput(input, self.target) #Backward
        if self.normalize:
          self.gradInput.div(torch.norm(self.gradInput, 1) + 1e-8) # Normalize Gradients
        self.gradInput_mul(self.strength)
        self.gradInput.add(gradOutput)
      else:
        self.target.resize_as_(gradOutput).copy_(gradOutput)
      return self.gradInput
        
class GramMatrix(nn.Module):
   def __init__(self, input):
       super(GramMatrix, self).__init__()

   def updateOutput(self, input):
     assert input.dim() == 3
     C, H, W = input.size(1), input.size(2), input.size(3)
     x_flat = input.view(C, H * W)
     self.output.resize(C, C)
     self.output.mm(x_flat, x_flat.t())
     return self.output

   def updateGradInput(self, input, gradOutput):
     assert input.dim() == 3 and input.size(1)
     C, H, W = input.size(1), input.size(2), input.size(3)
     x_flat = input.view(C, H * W)
     self.gradInput.resize(C, H * W).mm(gradOutput, x_flat)
     self.gradInput.addmm(gradOutput.t(), x_flat)
     self.gradInput = self.gradInput.view(C, H, W)
     return self.gradInput

# Define an nn Module to compute style loss in-place
class StyleLoss(nn.Module):
    def __init__(self, strength, normalize):
      super(StyleLoss, self).__init__()
      self.normalize = 'false'
      self.strength = strength
      self.target = torch.Tensor()
      self.mode = None
      self.loss = 0
    
      #self.gram = nn.GramMatrix()
      self.blend_weight = None
      self.G = None
      self.crit = nn.MSECriterion()
        
    def updateOutput(self, input):
      #self.G = self.gram.updateOutput(input) # Forward Gram
      self.G = GramMatrix(input).updateOutput # Forward Gram
      self.G.div(input.nelement())
      if self.mode == 'capture':
        if self.blend_weight == None:
          self.target.resize_as_(self.G).copy_(self.G)
        elif self.target.nelement() == 0:
          self.target.resize_as_(self.G).copy_(self.G).mul_(self.blend_weight)
        else:
          self.target.add(self.blend_weight, self.G)
      elif self.mode == 'loss':
        self.loss = self.strength * self.crit.updateOutput(input, self.target) #Forward
      self.output = input
      return self.output

    def updateGradInput(self, input, gradOutput):
      if self.mode == 'loss':
        dG = self.crit.updateGradInput(self.G, self.target) # Backward
        dG.div(input.nelement())
        self.gradInput = self.gram.updateGradInput(input) # Gram Backward
        if self.normalize:
          self.gradInput.div(torch.norm(self.gradInput, 1) + 1e-8) # Normalize Gradients
        self.gradInput_mul(self.strength)
        self.gradInput.add(gradOutput)
      else: 
        self.gradInput = gradOutput
      return self.gradInput

ModelSetup.py

import torch
import torch.nn as nn
import torch.legacy.nn as lnn
from torch.legacy.nn import SpatialConvolution
from torch.legacy.nn import SpatialMaxPooling
from torch.legacy.nn import ReLU
import torchvision
import torchvision.models as models
import torchvision.transforms as transforms

import copy

from LossModules import ContentLoss
from LossModules import StyleLoss
from LossModules import GramMatrix


use_cuda = torch.cuda.is_available()
dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor

num_channel1 = [3, 64, 64, 128, 128, 256, 256, 256, 256, 512, 512, 512, 512, 512, 512, 512, 512]
num_channel2 = [64, 64, 128, 128, 256, 256, 256, 256, 512, 512, 512, 512, 512, 512, 512, 512, 512]

def ModelSetup(cnn, style_image_caffe, content_image_caffe, style_weight, content_weight, content_layers, style_layers, normalize_gradients):
    # Fix the model layers
    print("Changing Layers")
    conv_num = [1]
    for idx, module in cnn.features._modules.items():
        if module.__class__.__name__ == 'ReLU':
            cnn.features._modules[idx] = lnn.ReLU()
        if module.__class__.__name__ == 'Conv2d':
            c = conv_num[0]
            channels_1 = num_channel1[c]
            channels_2 = num_channel2[c]
            cnn.features._modules[idx] = SpatialConvolution(channels_1, channels_2, 3, 3, 1, 1, 1, 1)
            conv_num[0] += 1
        if module.__class__.__name__ == 'MaxPool2d':
            cnn.features._modules[idx] = SpatialMaxPooling(2, 2, 2, 2, 0, 0)
    print("Layers Changed")
    cnn = cnn.features

   
    # Create the new network
    content_losses = []
    style_losses = []
    next_content_idx = 1
    next_style_idx = 1
    net = lnn.Sequential()
    net = net.cuda()

    print("-----------------------")
    for layer in list(cnn):
       print(layer)
    print("-----------------------")
    i = 1
    for layer in list(cnn):
      if next_content_idx <= len(content_layers) or next_style_idx <= len(style_layers):

        if isinstance(layer, lnn.SpatialConvolution):
            name = "conv_" + str(i)
            net.add(layer)

            if name in content_layers:
                print("Setting up content layer  "+ name)
                norm = normalize_gradients
                loss_module = ContentLoss(style_weight, norm)
                print(loss_module)
                net.add(loss_module)
                #content_losses.insert(next_content_idx, loss_module)
                #content_losses.insert(content_losses, loss_module)
                content_losses.append(loss_module)
                next_content_idx = next_content_idx + 1
            if name in style_layers:
                print("Setting up style layer  "+ name)
                norm = normalize_gradients
                loss_module = StyleLoss(style_weight, norm)#.type(dtype)
                net.add(loss_module)
                #style_losses.insert(style_losses, loss_module)
                #style_losses.insert(next_style_idx, loss_module)
                style_losses.append(loss_module)
                next_style_idx = next_style_idx + 1

        if isinstance(layer, lnn.ReLU):
            name = "relu_" + str(i)
            net.add(layer)

            if name in content_layers:
                # add content loss:
                target = net(content_image_caffe).clone()
                content_loss = ContentLoss(target, content_weight)
                net.add("content_loss_" + str(i), content_loss)
                content_losses.append(content_loss)

            if name in style_layers:
                # add style loss:
                target_feature = net(style_image_caffe).clone()
                target_feature_gram = GramMatrix(target_feature)
                style_loss = StyleLoss(target_feature_gram, style_weight)
                net.add("style_loss_" + str(i), style_loss)
                style_losses.append(style_loss)

        if isinstance(layer, lnn.SpatialMaxPooling):
            name = "pool_" + str(i)
            net.add(layer)  # ***

        i += 1

    return net, style_losses, content_losses

test2.py

# Code - Trying to translate https://github.com/jcjohnson/neural-style/blob/master/neural_style.lua to PyTorch.
from __future__ import print_function
import os
import sys

import torch
import torch.legacy.nn as nn
import torch.nn as nn2
from torch.autograd import Variable
import torch.legacy.optim as optim

from PIL import Image
from torch.legacy.nn import SpatialConvolution
from torch.legacy.nn import SpatialMaxPooling

import torchvision
import torchvision.transforms as transforms
import torchvision.models as models
from torchvision.utils import save_image

import copy

import argparse
parser = argparse.ArgumentParser()

# Basic options
parser.add_argument("-style_image", help="Style target image", default='examples/inputs/seated-nude.jpg')
parser.add_argument("-content_image", help="Content target image", default='examples/inputs/tubingen.jpg')
parser.add_argument("-image_size", help="Maximum height / width of generated image", type=int, default=512)

# Optimization options
parser.add_argument("-content_weight", help="content weight", type=int, default=5) 
parser.add_argument("-style_weight", help="style weight", type=int, default=10)

parser.add_argument("-num_iterations", help="iterations", type=int, default=1000)
parser.add_argument("-normalize_gradients", action='store_true')
parser.add_argument("-init", help="initialisation type", default="random", choices=["random", "image"])
parser.add_argument("-init_image", help="initial image", default="")
parser.add_argument("-optimizer", help="optimiser", default="lbfgs", choices=["lbfgs", "adam"])
parser.add_argument("-learning_rate", default=1)
parser.add_argument("-lbfgs_num_correction", help="lbfgs num correction", default=0)

# Output options  
parser.add_argument("-output_image", default='out.png')

# Other options
parser.add_argument("-style_scale", help="style scale", type=float, default=1.0)
#parser.add_argument("-proto_file", default='models/VGG_ILSVRC_19_layers_deploy.prototxt')
#parser.add_argument("-model_file", default='models/VGG_ILSVRC_19_layers.caffemodel')
parser.add_argument("-backend", choices=["nn", "cudnn", "clnn"], default='cudnn')
parser.add_argument("-seed", help="random number seed", default=-1)
params = parser.parse_args()


use_cuda = torch.cuda.is_available()
dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor

cnn = models.vgg19(pretrained=True)#.features

from ModelSetup import ModelSetup
from ImageSetup import SaveImage
from ImageSetup import ImageLoader

if use_cuda:
    cnn = cnn.cuda()

content_image_caffe = ImageLoader(params.content_image, params.image_size).type(dtype)
style_image_caffe = ImageLoader(params.style_image, params.image_size).type(dtype)


#content_layers_default = ['relu_4']
#style_layers_default = ['relu_1', 'relu_2', 'relu_3', 'relu_4', 'relu_5']
content_layers_default = ['conv_1']
style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']

def create_model(cnn, style_image_caffe, content_image_caffe, style_weight=params.style_weight, content_weight=params.style_weight, content_layers=content_layers_default, style_layers=style_layers_default, normalize_gradients=params.normalize_gradients):


    net, style_losses, content_losses = ModelSetup(cnn, style_image_caffe, content_image_caffe, style_weight=params.style_weight, content_weight=params.style_weight, content_layers=content_layers_default, style_layers=style_layers_default, normalize_gradients=params.normalize_gradients)
    return net, style_losses, content_losses


net, style_losses, content_losses = create_model(cnn, style_image_caffe, content_image_caffe, params.style_weight, params.content_weight, content_layers_default, style_layers_default)


print("Model Loaded")

# Capture content targets
for i in content_losses:
  i.mode = 'capture'
print("Capturing content targets")
#content_image_caffe = content_image_caffe.type(dtype)
#net.updateOutput(content_image_caffe.type(dtype))
#net.forward(content_image_caffe.type(dtype))
#content_image_caffe2 = ImageLoader(params.content_image, params.image_size)
net.forward(content_image_caffe.data.cpu())

# Capture style targets
for i in content_losses:
  i.mode = None

print("Capturing style target")
for j in style_losses:
  style_losses[j].mode = 'capture'
  style_losses[j].blend_weight = style_blend_weights[i]
net.updateOutput(style_image_caffe)


# Set all loss modules to loss mode
for i in content_losses:
  i.mode = 'loss'
for i in style_losses:
  i.mode = 'loss'

# Initialize the image
if params.seed >= 0:
  torch.manualSeed(params.seed)

# Run it through the network once to get the proper size for the gradient
# All the gradients will come from the extra loss modules, so we just pass
# zeros into the top of the net on the backward pass.
img = content_image_caffe.clone()
y = net.forward(img)
dy = img.clone().zero_() 


# Declaring this here lets us access it in maybe_print
optim_state = None
if params.optimizer == 'lbfgs':
  optim_state = {
    "maxIter": params.num_iterations,
    "verbose": True,
    "tolX":-1,
    "tolFun":-1,
  }
  if params.lbfgs_num_correction > 0:
    optim_state.nCorrection = params.lbfgs_num_correction
elif params.optimizer == 'adam':
    optim_state = {
      "learningRate": params.learning_rate,
    }

# Function to evaluate loss and gradient. We run the net forward and
# backward to get the gradient, and sum up losses from the loss modules.
# optim.lbfgs internally handles iteration and calls this function many
# times, so we manually count the number of iterations to handle printing
# and saving intermediate results.
num_calls = [0]

def feval(x):
  num_calls[0] += 1
  net.updateOutput(x)
  grad = net.updateGradInput(x, dy)
  loss = 0
  for n, mod in content_losses:
    loss = loss + mod.loss
  for n, mod in style_losses:
    loss = loss + mod.loss
  # optim.lbfgs expects a vector for gradients
  return loss, grad.view(grad.nelement())
    
# Run optimization.
if params.optimizer == 'lbfgs':
  print("Running optimization with L-BFGS")
  x, losses = optim.lbfgs(feval, img, optim_state)
elif params.optimizer == 'adam':
  print("Running optimization with ADAM")
  for t in params.num_iterations:
    x, losses = optim.adam(feval, img, optim_state)

save_image(output_img, params.output_image)