assignment3/StyleTransfer-PyTorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "pdf-title"
    ]
   },
   "source": [
    "# Style Transfer\n",
    "In this notebook we will implement the style transfer technique from [\"Image Style Transfer Using Convolutional Neural Networks\" (Gatys et al., CVPR 2015)](http://www.cv-foundation.org/openaccess/content_cvpr_2016/papers/Gatys_Image_Style_Transfer_CVPR_2016_paper.pdf).\n",
    "\n",
    "The general idea is to take two images, and produce a new image that reflects the content of one but the artistic \"style\" of the other. We will do this by first formulating a loss function that matches the content and style of each respective image in the feature space of a deep network, and then performing gradient descent on the pixels of the image itself.\n",
    "\n",
    "The deep network we use as a feature extractor is [SqueezeNet](https://arxiv.org/abs/1602.07360), a small model that has been trained on ImageNet. You could use any network, but we chose SqueezeNet here for its small size and efficiency.\n",
    "\n",
    "Here's an example of the images you'll be able to produce by the end of this notebook:\n",
    "\n",
    "![caption](example_styletransfer.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "pdf-ignore"
    ]
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "tags": [
     "pdf-ignore"
    ]
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.transforms as T\n",
    "import PIL\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from scipy.misc import imread\n",
    "from collections import namedtuple\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from cs231n.image_utils import SQUEEZENET_MEAN, SQUEEZENET_STD\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "pdf-ignore"
    ]
   },
   "source": [
    "We provide you with some helper functions to deal with images, since for this part of the assignment we're dealing with real JPEGs, not CIFAR-10 data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "tags": [
     "pdf-ignore"
    ]
   },
   "outputs": [],
   "source": [
    "def preprocess(img, size=512):\n",
    "    transform = T.Compose([\n",
    "        T.Resize(size),\n",
    "        T.ToTensor(),\n",
    "        T.Normalize(mean=SQUEEZENET_MEAN.tolist(),\n",
    "                    std=SQUEEZENET_STD.tolist()),\n",
    "        T.Lambda(lambda x: x[None]),\n",
    "    ])\n",
    "    return transform(img)\n",
    "\n",
    "def deprocess(img):\n",
    "    transform = T.Compose([\n",
    "        T.Lambda(lambda x: x[0]),\n",
    "        T.Normalize(mean=[0, 0, 0], std=[1.0 / s for s in SQUEEZENET_STD.tolist()]),\n",
    "        T.Normalize(mean=[-m for m in SQUEEZENET_MEAN.tolist()], std=[1, 1, 1]),\n",
    "        T.Lambda(rescale),\n",
    "        T.ToPILImage(),\n",
    "    ])\n",
    "    return transform(img)\n",
    "\n",
    "def rescale(x):\n",
    "    low, high = x.min(), x.max()\n",
    "    x_rescaled = (x - low) / (high - low)\n",
    "    return x_rescaled\n",
    "\n",
    "def rel_error(x,y):\n",
    "    return np.max(np.abs(x - y) / (np.maximum(1e-8, np.abs(x) + np.abs(y))))\n",
    "\n",
    "def features_from_img(imgpath, imgsize):\n",
    "    img = preprocess(PIL.Image.open(imgpath), size=imgsize)\n",
    "    img_var = img.type(dtype)\n",
    "    return extract_features(img_var, cnn), img_var\n",
    "\n",
    "# Older versions of scipy.misc.imresize yield different results\n",
    "# from newer versions, so we check to make sure scipy is up to date.\n",
    "def check_scipy():\n",
    "    import scipy\n",
    "    vnum = int(scipy.__version__.split('.')[1])\n",
    "    major_vnum = int(scipy.__version__.split('.')[0])\n",
    "    \n",
    "    assert vnum >= 16 or major_vnum >= 1, \"You must install SciPy >= 0.16.0 to complete this notebook.\"\n",
    "\n",
    "check_scipy()\n",
    "\n",
    "answers = dict(np.load('style-transfer-checks.npz'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "pdf-ignore"
    ]
   },
   "source": [
    "As in the last assignment, we need to set the dtype to select either the CPU or the GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "tags": [
     "pdf-ignore"
    ]
   },
   "outputs": [],
   "source": [
    "dtype = torch.FloatTensor\n",
    "# Uncomment out the following line if you're on a machine with a GPU set up for PyTorch!\n",
    "#dtype = torch.cuda.FloatTensor "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "tags": [
     "pdf-ignore"
    ]
   },
   "outputs": [],
   "source": [
    "# Load the pre-trained SqueezeNet model.\n",
    "cnn = torchvision.models.squeezenet1_1(pretrained=True).features\n",
    "cnn.type(dtype)\n",
    "\n",
    "# We don't want to train the model any further, so we don't want PyTorch to waste computation \n",
    "# computing gradients on parameters we're never going to update.\n",
    "for param in cnn.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "# We provide this helper code which takes an image, a model (cnn), and returns a list of\n",
    "# feature maps, one per layer.\n",
    "def extract_features(x, cnn):\n",
    "    \"\"\"\n",
    "    Use the CNN to extract features from the input image x.\n",
    "    \n",
    "    Inputs:\n",
    "    - x: A PyTorch Tensor of shape (N, C, H, W) holding a minibatch of images that\n",
    "      will be fed to the CNN.\n",
    "    - cnn: A PyTorch model that we will use to extract features.\n",
    "    \n",
    "    Returns:\n",
    "    - features: A list of feature for the input images x extracted using the cnn model.\n",
    "      features[i] is a PyTorch Tensor of shape (N, C_i, H_i, W_i); recall that features\n",
    "      from different layers of the network may have different numbers of channels (C_i) and\n",
    "      spatial dimensions (H_i, W_i).\n",
    "    \"\"\"\n",
    "    features = []\n",
    "    prev_feat = x\n",
    "    for i, module in enumerate(cnn._modules.values()):\n",
    "        next_feat = module(prev_feat)\n",
    "        features.append(next_feat)\n",
    "        prev_feat = next_feat\n",
    "    return features\n",
    "\n",
    "#please disregard warnings about initialization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing Loss\n",
    "\n",
    "We're going to compute the three components of our loss function now. The loss function is a weighted sum of three terms: content loss + style loss + total variation loss. You'll fill in the functions that compute these weighted terms below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "pdf-ignore"
    ]
   },
   "source": [
    "## Content loss\n",
    "We can generate an image that reflects the content of one image and the style of another by incorporating both in our loss function. We want to penalize deviations from the content of the content image and deviations from the style of the style image. We can then use this hybrid loss function to perform gradient descent **not on the parameters** of the model, but instead **on the pixel values** of our original image.\n",
    "\n",
    "Let's first write the content loss function. Content loss measures how much the feature map of the generated image differs from the feature map of the source image. We only care about the content representation of one layer of the network (say, layer $\\ell$), that has feature maps $A^\\ell \\in \\mathbb{R}^{1 \\times C_\\ell \\times H_\\ell \\times W_\\ell}$. $C_\\ell$ is the number of filters/channels in layer $\\ell$, $H_\\ell$ and $W_\\ell$ are the height and width. We will work with reshaped versions of these feature maps that combine all spatial positions into one dimension. Let $F^\\ell \\in \\mathbb{R}^{C_\\ell \\times M_\\ell}$ be the feature map for the current image and $P^\\ell \\in \\mathbb{R}^{C_\\ell \\times M_\\ell}$ be the feature map for the content source image where $M_\\ell=H_\\ell\\times W_\\ell$ is the number of elements in each feature map. Each row of $F^\\ell$ or $P^\\ell$ represents the vectorized activations of a particular filter, convolved over all positions of the image. Finally, let $w_c$ be the weight of the content loss term in the loss function.\n",
    "\n",
    "Then the content loss is given by:\n",
    "\n",
    "$L_c = w_c \\times \\sum_{i,j} (F_{ij}^{\\ell} - P_{ij}^{\\ell})^2$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def content_loss(content_weight, content_current, content_original):\n",
    "    \"\"\"\n",
    "    Compute the content loss for style transfer.\n",
    "    \n",
    "    Inputs:\n",
    "    - content_weight: Scalar giving the weighting for the content loss.\n",
    "    - content_current: features of the current image; this is a PyTorch Tensor of shape\n",
    "      (1, C_l, H_l, W_l).\n",
    "    - content_target: features of the content image, Tensor with shape (1, C_l, H_l, W_l).\n",
    "    \n",
    "    Returns:\n",
    "    - scalar content loss\n",
    "    \"\"\"\n",
    "    # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****\n\n",
    "    pass\n",
    "\n    # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test your content loss. You should see errors less than 0.001."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def content_loss_test(correct):\n",
    "    content_image = 'styles/tubingen.jpg'\n",
    "    image_size =  192\n",
    "    content_layer = 3\n",
    "    content_weight = 6e-2\n",
    "    \n",
    "    c_feats, content_img_var = features_from_img(content_image, image_size)\n",
    "    \n",
    "    bad_img = torch.zeros(*content_img_var.data.size()).type(dtype)\n",
    "    feats = extract_features(bad_img, cnn)\n",
    "    \n",
    "    student_output = content_loss(content_weight, c_feats[content_layer], feats[content_layer]).cpu().data.numpy()\n",
    "    error = rel_error(correct, student_output)\n",
    "    print('Maximum error is {:.3f}'.format(error))\n",
    "\n",
    "content_loss_test(answers['cl_out'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Style loss\n",
    "Now we can tackle the style loss. For a given layer $\\ell$, the style loss is defined as follows:\n",
    "\n",
    "First, compute the Gram matrix G which represents the correlations between the responses of each filter, where F is as above. The Gram matrix is an approximation to the covariance matrix -- we want the activation statistics of our generated image to match the activation statistics of our style image, and matching the (approximate) covariance is one way to do that. There are a variety of ways you could do this, but the Gram matrix is nice because it's easy to compute and in practice shows good results.\n",
    "\n",
    "Given a feature map $F^\\ell$ of shape $(C_\\ell, M_\\ell)$, the Gram matrix has shape $(C_\\ell, C_\\ell)$ and its elements are given by:\n",
    "\n",
    "$$G_{ij}^\\ell  = \\sum_k F^{\\ell}_{ik} F^{\\ell}_{jk}$$\n",
    "\n",
    "Assuming $G^\\ell$ is the Gram matrix from the feature map of the current image, $A^\\ell$ is the Gram Matrix from the feature map of the source style image, and $w_\\ell$ a scalar weight term, then the style loss for the layer $\\ell$ is simply the weighted Euclidean distance between the two Gram matrices:\n",
    "\n",
    "$$L_s^\\ell = w_\\ell \\sum_{i, j} \\left(G^\\ell_{ij} - A^\\ell_{ij}\\right)^2$$\n",
    "\n",
    "In practice we usually compute the style loss at a set of layers $\\mathcal{L}$ rather than just a single layer $\\ell$; then the total style loss is the sum of style losses at each layer:\n",
    "\n",
    "$$L_s = \\sum_{\\ell \\in \\mathcal{L}} L_s^\\ell$$\n",
    "\n",
    "Begin by implementing the Gram matrix computation below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def gram_matrix(features, normalize=True):\n",
    "    \"\"\"\n",
    "    Compute the Gram matrix from features.\n",
    "    \n",
    "    Inputs:\n",
    "    - features: PyTorch Tensor of shape (N, C, H, W) giving features for\n",
    "      a batch of N images.\n",
    "    - normalize: optional, whether to normalize the Gram matrix\n",
    "        If True, divide the Gram matrix by the number of neurons (H * W * C)\n",
    "    \n",
    "    Returns:\n",
    "    - gram: PyTorch Tensor of shape (N, C, C) giving the\n",
    "      (optionally normalized) Gram matrices for the N input images.\n",
    "    \"\"\"\n",
    "    # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****\n\n",
    "    pass\n",
    "\n    # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test your Gram matrix code. You should see errors less than 0.001."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gram_matrix_test(correct):\n",
    "    style_image = 'styles/starry_night.jpg'\n",
    "    style_size = 192\n",
    "    feats, _ = features_from_img(style_image, style_size)\n",
    "    student_output = gram_matrix(feats[5].clone()).cpu().data.numpy()\n",
    "    error = rel_error(correct, student_output)\n",
    "    print('Maximum error is {:.3f}'.format(error))\n",
    "    \n",
    "gram_matrix_test(answers['gm_out'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, implement the style loss:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Now put it together in the style_loss function...\n",
    "def style_loss(feats, style_layers, style_targets, style_weights):\n",
    "    \"\"\"\n",
    "    Computes the style loss at a set of layers.\n",
    "    \n",
    "    Inputs:\n",
    "    - feats: list of the features at every layer of the current image, as produced by\n",
    "      the extract_features function.\n",
    "    - style_layers: List of layer indices into feats giving the layers to include in the\n",
    "      style loss.\n",
    "    - style_targets: List of the same length as style_layers, where style_targets[i] is\n",
    "      a PyTorch Tensor giving the Gram matrix of the source style image computed at\n",
    "      layer style_layers[i].\n",
    "    - style_weights: List of the same length as style_layers, where style_weights[i]\n",
    "      is a scalar giving the weight for the style loss at layer style_layers[i].\n",
    "      \n",
    "    Returns:\n",
    "    - style_loss: A PyTorch Tensor holding a scalar giving the style loss.\n",
    "    \"\"\"\n",
    "    # Hint: you can do this with one for loop over the style layers, and should\n",
    "    # not be very much code (~5 lines). You will need to use your gram_matrix function.\n",
    "    # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****\n\n",
    "    pass\n",
    "\n    # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test your style loss implementation. The error should be less than 0.001."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def style_loss_test(correct):\n",
    "    content_image = 'styles/tubingen.jpg'\n",
    "    style_image = 'styles/starry_night.jpg'\n",
    "    image_size =  192\n",
    "    style_size = 192\n",
    "    style_layers = [1, 4, 6, 7]\n",
    "    style_weights = [300000, 1000, 15, 3]\n",
    "    \n",
    "    c_feats, _ = features_from_img(content_image, image_size)    \n",
    "    feats, _ = features_from_img(style_image, style_size)\n",
    "    style_targets = []\n",
    "    for idx in style_layers:\n",
    "        style_targets.append(gram_matrix(feats[idx].clone()))\n",
    "    \n",
    "    student_output = style_loss(c_feats, style_layers, style_targets, style_weights).cpu().data.numpy()\n",
    "    error = rel_error(correct, student_output)\n",
    "    print('Error is {:.3f}'.format(error))\n",
    "\n",
    "    \n",
    "style_loss_test(answers['sl_out'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Total-variation regularization\n",
    "It turns out that it's helpful to also encourage smoothness in the image. We can do this by adding another term to our loss that penalizes wiggles or \"total variation\" in the pixel values. \n",
    "\n",
    "You can compute the \"total variation\" as the sum of the squares of differences in the pixel values for all pairs of pixels that are next to each other (horizontally or vertically). Here we sum the total-variation regualarization for each of the 3 input channels (RGB), and weight the total summed loss by the total variation weight, $w_t$:\n",
    "\n",
    "$L_{tv} = w_t \\times \\left(\\sum_{c=1}^3\\sum_{i=1}^{H-1}\\sum_{j=1}^{W} (x_{i+1,j,c} - x_{i,j,c})^2 + \\sum_{c=1}^3\\sum_{i=1}^{H}\\sum_{j=1}^{W - 1} (x_{i,j+1,c} - x_{i,j,c})^2\\right)$\n",
    "\n",
    "In the next cell, fill in the definition for the TV loss term. To receive full credit, your implementation should not have any loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def tv_loss(img, tv_weight):\n",
    "    \"\"\"\n",
    "    Compute total variation loss.\n",
    "    \n",
    "    Inputs:\n",
    "    - img: PyTorch Variable of shape (1, 3, H, W) holding an input image.\n",
    "    - tv_weight: Scalar giving the weight w_t to use for the TV loss.\n",
    "    \n",
    "    Returns:\n",
    "    - loss: PyTorch Variable holding a scalar giving the total variation loss\n",
    "      for img weighted by tv_weight.\n",
    "    \"\"\"\n",
    "    # Your implementation should be vectorized and not require any loops!\n",
    "    # *****START OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****\n\n",
    "    pass\n",
    "\n    # *****END OF YOUR CODE (DO NOT DELETE/MODIFY THIS LINE)*****\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test your TV loss implementation. Error should be less  than 0.0001."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tv_loss_test(correct):\n",
    "    content_image = 'styles/tubingen.jpg'\n",
    "    image_size =  192\n",
    "    tv_weight = 2e-2\n",
    "\n",
    "    content_img = preprocess(PIL.Image.open(content_image), size=image_size).type(dtype)\n",
    "    \n",
    "    student_output = tv_loss(content_img, tv_weight).cpu().data.numpy()\n",
    "    error = rel_error(correct, student_output)\n",
    "    print('Error is {:.3f}'.format(error))\n",
    "    \n",
    "tv_loss_test(answers['tv_out'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we're ready to string it all together (you shouldn't have to modify this function):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false,
    "tags": [
     "pdf-ignore"
    ]
   },
   "outputs": [],
   "source": [
    "def style_transfer(content_image, style_image, image_size, style_size, content_layer, content_weight,\n",
    "                   style_layers, style_weights, tv_weight, init_random = False):\n",
    "    \"\"\"\n",
    "    Run style transfer!\n",
    "    \n",
    "    Inputs:\n",
    "    - content_image: filename of content image\n",
    "    - style_image: filename of style image\n",
    "    - image_size: size of smallest image dimension (used for content loss and generated image)\n",
    "    - style_size: size of smallest style image dimension\n",
    "    - content_layer: layer to use for content loss\n",
    "    - content_weight: weighting on content loss\n",
    "    - style_layers: list of layers to use for style loss\n",
    "    - style_weights: list of weights to use for each layer in style_layers\n",
    "    - tv_weight: weight of total variation regularization term\n",
    "    - init_random: initialize the starting image to uniform random noise\n",
    "    \"\"\"\n",
    "    \n",
    "    # Extract features for the content image\n",
    "    content_img = preprocess(PIL.Image.open(content_image), size=image_size).type(dtype)\n",
    "    feats = extract_features(content_img, cnn)\n",
    "    content_target = feats[content_layer].clone()\n",
    "\n",
    "    # Extract features for the style image\n",
    "    style_img = preprocess(PIL.Image.open(style_image), size=style_size).type(dtype)\n",
    "    feats = extract_features(style_img, cnn)\n",
    "    style_targets = []\n",
    "    for idx in style_layers:\n",
    "        style_targets.append(gram_matrix(feats[idx].clone()))\n",
    "\n",
    "    # Initialize output image to content image or nois\n",
    "    if init_random:\n",
    "        img = torch.Tensor(content_img.size()).uniform_(0, 1).type(dtype)\n",
    "    else:\n",
    "        img = content_img.clone().type(dtype)\n",
    "\n",
    "    # We do want the gradient computed on our image!\n",
    "    img.requires_grad_()\n",
    "    \n",
    "    # Set up optimization hyperparameters\n",
    "    initial_lr = 3.0\n",
    "    decayed_lr = 0.1\n",
    "    decay_lr_at = 180\n",
    "\n",
    "    # Note that we are optimizing the pixel values of the image by passing\n",
    "    # in the img Torch tensor, whose requires_grad flag is set to True\n",
    "    optimizer = torch.optim.Adam([img], lr=initial_lr)\n",
    "    \n",
    "    f, axarr = plt.subplots(1,2)\n",
    "    axarr[0].axis('off')\n",
    "    axarr[1].axis('off')\n",
    "    axarr[0].set_title('Content Source Img.')\n",
    "    axarr[1].set_title('Style Source Img.')\n",
    "    axarr[0].imshow(deprocess(content_img.cpu()))\n",
    "    axarr[1].imshow(deprocess(style_img.cpu()))\n",
    "    plt.show()\n",
    "    plt.figure()\n",
    "    \n",
    "    for t in range(200):\n",
    "        if t < 190:\n",
    "            img.data.clamp_(-1.5, 1.5)\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        feats = extract_features(img, cnn)\n",
    "        \n",
    "        # Compute loss\n",
    "        c_loss = content_loss(content_weight, feats[content_layer], content_target)\n",
    "        s_loss = style_loss(feats, style_layers, style_targets, style_weights)\n",
    "        t_loss = tv_loss(img, tv_weight) \n",
    "        loss = c_loss + s_loss + t_loss\n",
    "        \n",
    "        loss.backward()\n",
    "\n",
    "        # Perform gradient descents on our image values\n",
    "        if t == decay_lr_at:\n",
    "            optimizer = torch.optim.Adam([img], lr=decayed_lr)\n",
    "        optimizer.step()\n",
    "\n",
    "        if t % 100 == 0:\n",
    "            print('Iteration {}'.format(t))\n",
    "            plt.axis('off')\n",
    "            plt.imshow(deprocess(img.data.cpu()))\n",
    "            plt.show()\n",
    "    print('Iteration {}'.format(t))\n",
    "    plt.axis('off')\n",
    "    plt.imshow(deprocess(img.data.cpu()))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate some pretty pictures!\n",
    "\n",
    "Try out `style_transfer` on the three different parameter sets below. Make sure to run all three cells. Feel free to add your own, but make sure to include the results of style transfer on the third parameter set (starry night) in your submitted notebook.\n",
    "\n",
    "* The `content_image` is the filename of content image.\n",
    "* The `style_image` is the filename of style image.\n",
    "* The `image_size` is the size of smallest image dimension of the content image (used for content loss and generated image).\n",
    "* The `style_size` is the size of smallest style image dimension.\n",
    "* The `content_layer` specifies which layer to use for content loss.\n",
    "* The `content_weight` gives weighting on content loss in the overall loss function. Increasing the value of this parameter will make the final image look more realistic (closer to the original content).\n",
    "* `style_layers` specifies a list of which layers to use for style loss. \n",
    "* `style_weights` specifies a list of weights to use for each layer in style_layers (each of which will contribute a term to the overall style loss). We generally use higher weights for the earlier style layers because they describe more local/smaller scale features, which are more important to texture than features over larger receptive fields. In general, increasing these weights will make the resulting image look less like the original content and more distorted towards the appearance of the style image.\n",
    "* `tv_weight` specifies the weighting of total variation regularization in the overall loss function. Increasing this value makes the resulting image look smoother and less jagged, at the cost of lower fidelity to style and content. \n",
    "\n",
    "Below the next three cells of code (in which you shouldn't change the hyperparameters), feel free to copy and paste the parameters to play around them and see how the resulting image changes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Composition VII + Tubingen\n",
    "params1 = {\n",
    "    'content_image' : 'styles/tubingen.jpg',\n",
    "    'style_image' : 'styles/composition_vii.jpg',\n",
    "    'image_size' : 192,\n",
    "    'style_size' : 512,\n",
    "    'content_layer' : 3,\n",
    "    'content_weight' : 5e-2, \n",
    "    'style_layers' : (1, 4, 6, 7),\n",
    "    'style_weights' : (20000, 500, 12, 1),\n",
    "    'tv_weight' : 5e-2\n",
    "}\n",
    "\n",
    "style_transfer(**params1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scream + Tubingen\n",
    "params2 = {\n",
    "    'content_image':'styles/tubingen.jpg',\n",
    "    'style_image':'styles/the_scream.jpg',\n",
    "    'image_size':192,\n",
    "    'style_size':224,\n",
    "    'content_layer':3,\n",
    "    'content_weight':3e-2,\n",
    "    'style_layers':[1, 4, 6, 7],\n",
    "    'style_weights':[200000, 800, 12, 1],\n",
    "    'tv_weight':2e-2\n",
    "}\n",
    "\n",
    "style_transfer(**params2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Starry Night + Tubingen\n",
    "params3 = {\n",
    "    'content_image' : 'styles/tubingen.jpg',\n",
    "    'style_image' : 'styles/starry_night.jpg',\n",
    "    'image_size' : 192,\n",
    "    'style_size' : 192,\n",
    "    'content_layer' : 3,\n",
    "    'content_weight' : 6e-2,\n",
    "    'style_layers' : [1, 4, 6, 7],\n",
    "    'style_weights' : [300000, 1000, 15, 3],\n",
    "    'tv_weight' : 2e-2\n",
    "}\n",
    "\n",
    "style_transfer(**params3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "## Feature Inversion\n",
    "\n",
    "The code you've written can do another cool thing. In an attempt to understand the types of features that convolutional networks learn to recognize, a recent paper [1] attempts to reconstruct an image from its feature representation. We can easily implement this idea using image gradients from the pretrained network, which is exactly what we did above (but with two different feature representations).\n",
    "\n",
    "Now, if you set the style weights to all be 0 and initialize the starting image to random noise instead of the content source image, you'll reconstruct an image from the feature representation of the content source image. You're starting with total noise, but you should end up with something that looks quite a bit like your original image.\n",
    "\n",
    "(Similarly, you could do \"texture synthesis\" from scratch if you set the content weight to 0 and initialize the starting image to random noise, but we won't ask you to do that here.) \n",
    "\n",
    "Run the following cell to try out feature inversion.\n",
    "\n",
    "[1] Aravindh Mahendran, Andrea Vedaldi, \"Understanding Deep Image Representations by Inverting them\", CVPR 2015\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Feature Inversion -- Starry Night + Tubingen\n",
    "params_inv = {\n",
    "    'content_image' : 'styles/tubingen.jpg',\n",
    "    'style_image' : 'styles/starry_night.jpg',\n",
    "    'image_size' : 192,\n",
    "    'style_size' : 192,\n",
    "    'content_layer' : 3,\n",
    "    'content_weight' : 6e-2,\n",
    "    'style_layers' : [1, 4, 6, 7],\n",
    "    'style_weights' : [0, 0, 0, 0], # we discard any contributions from style to the loss\n",
    "    'tv_weight' : 2e-2,\n",
    "    'init_random': True # we want to initialize our image to be random\n",
    "}\n",
    "\n",
    "style_transfer(**params_inv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}