Pytorch Implementation of Perceptual Losses for Real-Time Style Transfer | by Ceshine Lee | Towards Data Science

In this post I’ll briefly go through my experience of coding and training real-time style transfer models in Pytorch. The work is heavily based on Abhishek Kadian’s implementation, which works perfectly Fine. I’ve made some modification both for fun and to be more familiar with Pytorch.
The model uses the method described in Perceptual Losses for Real-Time Style Transfer and Super-Resolution along with Instance Normalization. (Super-Resolution is not implemented)
Three major parts I’ve added to the implementation:
Use the official pre-trained VGG modelOutput intermediate results during trainingAdd Total Variation Regularization as described in the paper
Use the official pre-trained VGG model
Output intermediate results during training
Add Total Variation Regularization as described in the paper
First we need to take a quick look at the model structure. The main contribution of the paper is proposing that feeding forward the generated image to a pre-trained image classification model and extract the output from some intermediate layers to calculate losses would produce similar results of Gatys et al but with significantly less computational resources. So the first part of the structure is a “Image Transform Net” which generate new image from the input image. And the second part is simply a “Loss Network”, which is the feeding forward part.The weight of the loss network is fixed and will not be updated during training.
Abhishek’s implementation uses a traditional VGG model with BGR channel order and [-103.939, -116.779, -123.680] offsets to center channel means (it seems to also be what the paper used). The official pytorch pre-trained models use a unified format:
All pre-trained models expect input images normalized in the same way, i.e. mini-batches of 3-channel RGB images of shape (3 x H x W), where H and W are expected to be at least 224. The images have to be loaded in to a range of [0, 1] and then normalized using mean = [0.485, 0.456, 0.406] and std = [0.229, 0.224, 0.225]
Here’s the code to extract outputs from the official pre-trained model:
LossOutput = namedtuple("LossOutput", ["relu1_2", "relu2_2", "relu3_3", "relu4_3"])# https://discuss.pytorch.org/t/how-to-extract-features-of-an-image-from-a-trained-model/119/3class LossNetwork(torch.nn.Module):    def __init__(self, vgg_model):        super(LossNetwork, self).__init__()        self.vgg_layers = vgg_model.features        self.layer_name_mapping = {            '3': "relu1_2",            '8': "relu2_2",            '15': "relu3_3",            '22': "relu4_3"        }        def forward(self, x):        output = {}        for name, module in self.vgg_layers._modules.items():            x = module(x)            if name in self.layer_name_mapping:                output[self.layer_name_mapping[name]] = x        return LossOutput(**output)
To initiate:
vgg_model = vgg.vgg16(pretrained=True)if torch.cuda.is_available():    vgg_model.cuda()loss_network = LossNetwork(vgg_model)loss_network.eval()
There’s no batch normalization in VGG model unless explicitly specified. So the values of activation is on a significant different scale from the previous implementation. Generally you’d need to scale up the style loss (gram matrix) because most activation is less than 1, and taking dot products would make it smaller.
This helps a lot when tuning content weight against style weight ratio. You can stop training and re-tune the parameters during training, instead of having to wait 4 hours for the training to complete.
The paper mentioned this in the Experiment section, but it seems Abhishek didn’t implement it:
The output images are regularized with total variation regularization with a strength of between 1 × 10e-6 and 1 ×10e-4, chosen via cross-validation per style target.
It is quite easy to implement:
reg_loss = REGULARIZATION * (    torch.sum(torch.abs(y[:, :, :, :-1] - y[:, :, :, 1:])) +     torch.sum(torch.abs(y[:, :, :-1, :] - y[:, :, 1:, :])))
Pytorch autograd will handle backward propagation for you. In practice I haven’t found out how to tune the regularization weight properly. So far the weights I used seemed not making much difference in the output images.
The model is trained using Microsoft COCO dataset. The image is resized to 256 x 256, and network is trained around 2 epochs with a batch size of 4 (the same as the paper). The training time using a GTX1070 is around 4 to 4.5 hours, on par with what the paper reports. Based on my rough experiment, a lot of computation time is used in normalizing the input images. The training might be faster if we use the original VGG model (not tested). The content weight v.s. style ratio is usually set to 1 : 10e3 ~ 10e5 after some manual tuning.
Because the network is fully convolutional, you could feed much bigger or smaller images than 256 x 256 to the network in test time. I write some script to transform animated GIFs and videos using scikit-video and ffmpeg for fun:
Original video credit: Backpacking Around New Zealand. The processing speed is around 4 frames per second. I didn’t align the frame rate, so the time in the following videos is slower than the original one.
Remember to clip output numpy array to [0, 255] range and cast it to uint8. Otherwise matplot.pyplot.imshow would show weird results. Spent quite some time debugging this in the wrong direction thinking the the model training code has some serious bugs.Remember to use model.train() and model.eval(). It only has effects on models containing dropout or batch normalization layers, but it is a good habit to keep. Especially easy to forget if you’re coming from Keras.
Remember to clip output numpy array to [0, 255] range and cast it to uint8. Otherwise matplot.pyplot.imshow would show weird results. Spent quite some time debugging this in the wrong direction thinking the the model training code has some serious bugs.
Remember to use model.train() and model.eval(). It only has effects on models containing dropout or batch normalization layers, but it is a good habit to keep. Especially easy to forget if you’re coming from Keras.
Sometimes the network generate some weird patches in open areas. Don’t know where it’s from and how to fix it yet.Maybe assigning different weights to relu1_2, relu2_2, relu3_3, relu4_3 outputs for different style image would yield better results.Try using different pre-trained networks as the loss network.Try implement Super-Resolution as wellCompare results from CycleGAN. I tried CycleGAN but the training time is too long and I lost patience at the time. Should try again later.Regularize against variation between consecutive frames when generating videos. This should help reduce flickering during playback. I’ve heard people using this technique, but have no idea how to actually do it yet.
Sometimes the network generate some weird patches in open areas. Don’t know where it’s from and how to fix it yet.
Maybe assigning different weights to relu1_2, relu2_2, relu3_3, relu4_3 outputs for different style image would yield better results.
Try using different pre-trained networks as the loss network.
Try implement Super-Resolution as well
Compare results from CycleGAN. I tried CycleGAN but the training time is too long and I lost patience at the time. Should try again later.
Regularize against variation between consecutive frames when generating videos. This should help reduce flickering during playback. I’ve heard people using this technique, but have no idea how to actually do it yet.
The code is located at this Github repo. Since this is basically a personal fun project, the documentation is nonexistent and the main part of the code resides in two Jupyter notebooks: style-transfer.ipynb and Video.ipynb. Sorry about that.
(2011/11/30 Update: Fixed links to the notebook above. In the meantime, a newer version of the notebooks is available here: 01-image-style-transfer.ipynb and 03-Video-Creation.ipynb.)