Neural_Style_Transfer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Imports"
      ],
      "metadata": {
        "id": "lmW87YgDFYEA"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "Dnj6nbEIDORy"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from PIL import Image\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.models as models\n",
        "from torchvision.utils import save_image"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Modified VGG19 for Neural Style Transfer "
      ],
      "metadata": {
        "id": "b8yDISAgFa1n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class VGG(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(VGG, self).__init__()\n",
        "\n",
        "        self.chosen_features = [\"0\", \"5\", \"10\", \"19\", \"28\"]\n",
        "        self.model = models.vgg19(pretrained=True).features[:29]\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Store relevant features\n",
        "        features = []\n",
        "        for layer_num, layer in enumerate(self.model):\n",
        "            x = layer(x)\n",
        "\n",
        "            if str(layer_num) in self.chosen_features:\n",
        "                features.append(x)\n",
        "\n",
        "        return features\n"
      ],
      "metadata": {
        "id": "Wd-xnEg4DkWB"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def load_image(image_name):\n",
        "    image = Image.open(image_name)\n",
        "    image = loader(image).unsqueeze(0)\n",
        "    return image.to(device)"
      ],
      "metadata": {
        "id": "ZrdkK7n-Dl_C"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "imsize = 356"
      ],
      "metadata": {
        "id": "mQCdPPgWDpBR"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "loader = transforms.Compose(\n",
        "    [\n",
        "        transforms.Resize((imsize, imsize)),\n",
        "        transforms.ToTensor(),\n",
        "    ]\n",
        ")"
      ],
      "metadata": {
        "id": "fDLtQ5GgDtSZ"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "original_img = load_image(\"/content/annahathaway.png\")\n",
        "style_img = load_image(\"/content/style.jpg\")"
      ],
      "metadata": {
        "id": "c8LiG8qFDvUI"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "generated = original_img.clone().requires_grad_(True)\n",
        "model = VGG().to(device).eval()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6FZ4ssoLDxO-",
        "outputId": "0b2afe3f-15f4-45df-d22a-86842d2005ba"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.10/dist-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG19_Weights.IMAGENET1K_V1`. You can also use `weights=VGG19_Weights.DEFAULT` to get the most up-to-date weights.\n",
            "  warnings.warn(msg)\n",
            "Downloading: \"https://download.pytorch.org/models/vgg19-dcbb9e9d.pth\" to /root/.cache/torch/hub/checkpoints/vgg19-dcbb9e9d.pth\n",
            "100%|██████████| 548M/548M [00:01<00:00, 298MB/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "total_steps = 6000\n",
        "learning_rate = 0.001\n",
        "alpha = 1\n",
        "beta = 0.01\n",
        "optimizer = optim.Adam([generated], lr=learning_rate)"
      ],
      "metadata": {
        "id": "9f-3JsojD1e7"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for step in range(total_steps):\n",
        "\n",
        "    generated_features = model(generated)\n",
        "    original_img_features = model(original_img)\n",
        "    style_features = model(style_img)\n",
        "\n",
        "    style_loss = 0\n",
        "    original_loss = 0\n",
        "\n",
        "    for gen_feature, orig_feature, style_feature in zip(\n",
        "        generated_features, original_img_features, style_features\n",
        "    ):\n",
        "\n",
        "        batch_size, channel, height, width = gen_feature.shape\n",
        "        original_loss += torch.mean((gen_feature - orig_feature) ** 2)\n",
        "        G = gen_feature.view(channel, height * width).mm(\n",
        "            gen_feature.view(channel, height * width).t()\n",
        "        )\n",
        "        A = style_feature.view(channel, height * width).mm(\n",
        "            style_feature.view(channel, height * width).t()\n",
        "        )\n",
        "        style_loss += torch.mean((G - A) ** 2)\n",
        "\n",
        "    total_loss = alpha * original_loss + beta * style_loss\n",
        "    optimizer.zero_grad()\n",
        "    total_loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "    if step % 200 == 0:\n",
        "        print(total_loss)\n",
        "        save_image(generated, \"generated.png\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ECbsy1-jD4xU",
        "outputId": "30d0e477-2fed-4469-91ae-d18f767032b8"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor(1579634.6250, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(84241.6094, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(39178.7891, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(28035.4316, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(22958.8223, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(19641.7637, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(17160.8711, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(15153.3203, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(13450.8691, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(11972.6250, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(10683.7480, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(9567.0312, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(8606.8037, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(7783.1680, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(7080.0557, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(6476.3057, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(5955.1548, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(5501.7974, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(5103.8979, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(4757.0063, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(4451.4326, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(4181.9028, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(3942.8213, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(3729.6304, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(3539.8633, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(3372.7090, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(3223.6177, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(3088.6531, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(2969.1299, device='cuda:0', grad_fn=<AddBackward0>)\n",
            "tensor(2862.4849, device='cuda:0', grad_fn=<AddBackward0>)\n"
          ]
        }
      ]
    }
  ]
}