main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Code Authors\n",
    "Özgür Aslan 2236958 aslan.ozgur@metu.edu.tr\n",
    "Burak Bolat 2237097 burak.bolat@metu.edu.tr\n",
    "\n",
    "### Paper Information\n",
    "The paper we selected to implement is [HistoGAN: Controlling Colors of GAN-Generated and Real Images via Color Histograms](https://arxiv.org/abs/2011.11731)\n",
    "The main idea of the paper is to use color histogram of target images to control the colors of the generated image without changing the high level features of the generated image (gender, having glasses, beard, hair style and objects in the background...)  \n",
    "To accomplish this idea, they modify the StyleGAN2 architecture:\n",
    "- In the last 2 style blocks, instead of using affine transformation of the w vector, they use the color histogram projected by a neural network.\n",
    "- Different from mixing regularization, they use a histogram based loss.\n",
    "- The histogram loss uses two different target images to compute color histograms and interpolate this histograms to obtain a new one. The interpolated histogram is given to generator network to generate a target image with colors controlled by the interpolated histogram. This way the authors try to prevent the generator network to overfit color histograms of the trained dataset.\n",
    "- Due to hardware limitations the network does not generate 1024x1024 resolution images but generates 256x256 images.\n",
    "- Also due to hardware limitations they use batches of size 2 with gradient accumulation.\n",
    "\n",
    "![arch](materials/arch.png)\n",
    "\n",
    "#### Histogram Computation\n",
    "Computing Histogram is critical since it directly affects style of last 2 blocks. Authors used chrominance logarithm space. It normalizes each color channel with respect to other two channels in logarithmic space. In this chrominance space, there is u and v axes. That is, if we look at red channel's chrominance space, u is the normalization of red channel with respect to green and v is the normalization of red channel with respect to blue. Same holds for all color channels.  \n",
    "\n",
    "After shifting RGB space to RGB-uv space, the histogram is computed as it is computationally efficient and more stable. Authors used 64 bin for the histogram which results in 64x64 histogram for u and v channel. We have 3 channels, namely red, green and blue, thus, overall the histogram is 3x64x64. Histogram is weighted with respect to pixel intesity, i.e. if a pixel has high RGB values its affect on the histogram bin is higher. Last difference of the histogram than histograms of previous works is kernels for computing bins. Authors do not used exact bin selection. Instead, they put a normalized pixel into a bin with respect to soft kernel. That means, if we have a red channel after normalized with respect to green and blue, we have some u and v values. Instead of just adding 1 (1 being chosen for simplicty, remember intensity multiplication) to the bin of H(u,v), they add values to the neighbour of (u,v) with the value after inverse quadratic kernel.\n",
    "\n",
    "Histogram feature is computed like syle vector (w). It passed through the same neural network architecture with different parametes. More precisely, histogram passes through 8 layer MLP and outputs latent histogram vector size of 512.  \n",
    "\n",
    "We put some computed histograms by us. Images taken from internet crawling.  \n",
    "\n",
    "![asd](materials/gresized.png) ![asc](materials/ghist1.png)  \n",
    "![asb](materials/rresized.png) ![asj](materials/rhist1.png)\n",
    "\n",
    "#### Loss for Training with Histogram\n",
    "Since the paper uses target histogram for generation, generated image should have close histogram to target. Thus a closeness measure Hellinger distance between histogram of generated and target images is computed and tried to minimize. You can check the losses belove.\n",
    "\n",
    "Difference between histograms  \n",
    "![l1](materials/hloss.png)\n",
    "\n",
    "Total loss for generator  \n",
    "![lt](materials/total_loss.png)\n",
    "\n",
    "### Discriminator\n",
    "\n",
    "Discriminator consist of residual blocks. There are log_2(N)-1 such bloks where N is image resolution, to be spesific 256. As a result, the discriminator has 7 layers. First block takes 3 channel image as input and outputs m channel features. After the first block, each block produces 2*m of previous block. At the end of residual blocks, a FC layer outputs a scaler.\n",
    "\n",
    "![res](materials/residual.png)\n",
    "\n",
    "### Important Note on Dataset\n",
    "The Anime Face Dataset is a Kaggle dataset, thus, requires a Kaggle account. Using an account one can download it from:\n",
    "https://www.kaggle.com/datasets/splcher/animefacedataset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Faced Challenges\n",
    "\n",
    "#### Architecture Challenges\n",
    "\n",
    "As we stated, HistoGAN is built on StyleGAN2. StyleGAN2 scales directly the weigth of the model, unlike the first version does the nearly same operations, namely mod-demod, on convolved image (or say not directly weights for convoltion filters). The original StyleGAN2 implemented using Tensorflow which allows to multiplication on weights, that is called in-place operation on variables. However, Pytorch does not allow in-place operations on built in modules like torch.nn.Conv2d. Therefore, we implemented a conv2d version. Model parameters are Pytorch Variables and convolution operation is handled with fold and unfold operations of Pytorch. Doing so we can apply convolution after scaling weights of convolution filters.\n",
    "\n",
    "#### Saturation of Generator\n",
    "\n",
    "After implementing StyleGAN2 and HistoGAN, we tried to train models. We saw that generator of HistoGAN does not learn and tried to train StyleGAN (with the shallow version that HistoGAN uses, it produces 256x256 images). However, we get rapid saturation of generator and have not solved yet. Here we present some generated images from training. \n",
    "\n",
    "![m1](materials/fake_0_199_0.png) ![m2](materials/fake_0_199_1.png) ![m3](materials/fake_0_399_0.png) ![m4](materials/fake_0_599_0.png) \n",
    "\n",
    "#### Training Challenges\n",
    "\n",
    "During the training phase, the paper does not mentioned how generator outputs the images. We made different assumption such as using sigmoid or tanh to generate pixels in a range. Another assumption for the same problem is using ReLU or leaky ReLU that we saw from other generator implementation.  \n",
    "StyleGAN2 stated that they used non saturating loss for some datasets and WGAN-GP loss for other datasets. HistoGAN paper does not clearly mention on this. Consequently, we implemented both but non saturating loss lead numerical issues like nan or infs. On the other hand, WGAN-GP computes high loss values and results in rapid saturation (see above figures). This issues may be resulted from hand implemented convolution operations. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Implementation Info\n",
    "- python 3.7.13\n",
    "- pytorch 1.11.0 with cuda10.2   \n",
    "We used conda environments for clean library setups and included environment.yml file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import torch\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import transforms\n",
    "from torchvision.utils import save_image\n",
    "\n",
    "from data import AnimeFacesDataset\n",
    "from model import Discriminator, HistoGAN\n",
    "from loss import compute_gradient_penalty, pl_reg, r1_reg, wgan_gp_disc_loss, wgan_gp_gen_loss\n",
    "from utils import random_interpolate_hists, histogram_feature_v2\n",
    "\n",
    "import os\n",
    "\n",
    "# for debugging\n",
    "torch.autograd.set_detect_anomaly(True)\n",
    "# for faster trainingg\n",
    "torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "real_image_dir = \"images\"\n",
    "transform = transforms.Compose(\n",
    "        [transforms.Resize((256,256)),\n",
    "        transforms.RandomHorizontalFlip(0.5)])\n",
    "dataset = AnimeFacesDataset(real_image_dir, transform, device)\n",
    "# due to hardware limitations similar to paper's authors we kept the batch size small\n",
    "batch_size = 2\n",
    "# the dataset contains 63,632 datum, and the we could not make the network to generate meaningfull images therefore kept the epochs small and experimented\n",
    "num_epochs = 2\n",
    "# variable to hold after how many discriminator updates to update the generator\n",
    "g_update_iter = 5\n",
    "# after how many gradient accumulation to optimize parameters\n",
    "acc_gradient_iter = 1\n",
    "# scalar of R1 regularization\n",
    "r1_factor = 10\n",
    "# variables for Path length regularization\n",
    "# please see StyleGAN2 paper B. Implementation Details Path length regularization\n",
    "ema_decay_coeff = 0.99\n",
    "target_scale = torch.tensor([0]).to(device)\n",
    "plr_factor = np.log(2)/(256**2*(np.log(256)-np.log(2)))\n",
    "# after how many iterations to save the nework parameters and generated images\n",
    "save_iter = 200\n",
    "# path to save generated images\n",
    "fake_image_dir = \"generated_images\"\n",
    "if not os.path.isdir(fake_image_dir):\n",
    "    os.mkdir(fake_image_dir)\n",
    "# number of residual blocks in the discriminator \n",
    "num_res_blocks = 7\n",
    "# network capacity to decide the intermediate channel sizes of discrimimator and learnable constant channel size of generator \n",
    "network_capacity = 16 \n",
    "# histogram's bin size\n",
    "bin_size = 64\n",
    "# the number of channels are decides as log2(image_res) -1 since we generate 256 res images, there are 7 channels\n",
    "generator_channel_sizes = [1024, 512, 512, 512, 256, 128, 64]\n",
    "learning_rate = 2e-4\n",
    "# coefficient of gradient penalty\n",
    "coeff_penalty = 10 # same as the StyleGAN2 paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Dataset, Dataloader, Discriminator and Generator\n",
    "dataset = AnimeFacesDataset(real_image_dir, transform, device)\n",
    "dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "generator = HistoGAN(network_capacity, bin_size, generator_channel_sizes)\n",
    "discriminator = Discriminator(num_res_blocks, network_capacity)\n",
    "\n",
    "# If a pretrained network exists, load their parameters to continue training\n",
    "if os.path.exists(\"generator.pt\"):\n",
    "    generator.load_state_dict(torch.load(\"generator.pt\"))\n",
    "if os.path.exists(\"discriminator.pt\"):\n",
    "    discriminator.load_state_dict(torch.load(\"discriminator.pt\"))\n",
    "\n",
    "\n",
    "discriminator = discriminator.to(device)\n",
    "generator=generator.to(device)\n",
    "\n",
    "# Initialize optimizers \n",
    "gene_optim = torch.optim.Adam(generator.parameters(), lr= learning_rate)\n",
    "disc_optim = torch.optim.Adam(discriminator.parameters(), lr= learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'num_epochs' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_25697/3708850999.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# Traning loop without gradient accumulation\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;31m# Gradient accumulation is implemented and tried in train2.py but had some performance and memory consumption issues therefore not added here\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnum_epochs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0miter\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_data\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataloader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0;31m# torch.cuda.empty_cache()\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'num_epochs' is not defined"
     ]
    }
   ],
   "source": [
    "# Traning loop without gradient accumulation\n",
    "# Gradient accumulation is implemented and tried in train2.py but had some performance and memory consumption issues therefore not added here\n",
    "for epoch in range(num_epochs):\n",
    "    for iter, batch_data in enumerate(dataloader):\n",
    "        # torch.cuda.empty_cache() \n",
    "        training_percent = 100*iter*batch_data.size(0)/len(dataset)\n",
    "        batch_data = batch_data.to(device)\n",
    "        # Sample random Gaussian noise\n",
    "        z = torch.randn(batch_data.size(0), 512).to(device)\n",
    "        # Interpolate between target image histogram \n",
    "        # to prevent overfitting to dataset images\n",
    "        target_hist = random_interpolate_hists(batch_data)\n",
    "        # Generate fake images\n",
    "        fake_data, w = generator(z, target_hist)\n",
    "\n",
    "        # Detach fake data so no gradient accumalition \n",
    "        # to generator while only training discriminator\n",
    "        fake_data = fake_data.detach()\n",
    "\n",
    "        # Compute real probabilities computed by discriminator\n",
    "        fake_scores = discriminator(fake_data)\n",
    "        real_scores = discriminator(batch_data)\n",
    "        gradient_penalty = compute_gradient_penalty(fake_data, batch_data, discriminator)\n",
    "        d_loss = wgan_gp_disc_loss(real_scores, fake_scores, gradient_penalty, coeff_penalty)\n",
    "        #d_loss = disc_loss(fake_scores, real_scores)\n",
    "        # in stylegan2 paper they argue applying regularization in every 16 iteration does not hurt perfrormance \n",
    "        if (iter+1) % 16 == 0: \n",
    "            # r1 regulatization\n",
    "            d_loss = d_loss + r1_reg(batch_data, discriminator, r1_factor)  \n",
    "\n",
    "        print(\"%\", training_percent, \" Disc loss:\", d_loss.item())\n",
    "        d_loss.backward()\n",
    "        disc_optim.step()\n",
    "        disc_optim.zero_grad()\n",
    "\n",
    "        if (iter+1) % g_update_iter == 0:\n",
    "            z = torch.randn(batch_data.size(0), 512).to(device)\n",
    "            fake_data, w = generator(z, target_hist) \n",
    "\n",
    "            disc_score = discriminator(fake_data)\n",
    "            g_loss = wgan_gp_gen_loss(disc_score)\n",
    "            if (iter+1) % (8*g_update_iter) == 0:\n",
    "                plr, ema_decay_coeff = pl_reg(fake_data, w, target_scale, plr_factor, ema_decay_coeff)\n",
    "                g_loss = g_loss + plr\n",
    "\n",
    "            print(\"%\", training_percent, \"Gen loss:\", g_loss.item())\n",
    "            g_loss.backward()\n",
    "            gene_optim.step()\n",
    "            gene_optim.zero_grad()\n",
    "            \n",
    "        if (iter+1) % save_iter == 0:\n",
    "            for i in range(fake_data.size(0)):\n",
    "                save_image(fake_data[i], os.path.join(fake_image_dir, \"fake_{}_{}_{}.png\".format(epoch, iter, i)))\n",
    "            torch.save(generator.state_dict(), \"generator.pt\")\n",
    "            torch.save(discriminator.state_dict(), \"discriminator.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "45dd755165ed25490c252c1c43843a664f052864db6e3554fdaf500ce1613bda"
  },
  "kernelspec": {
   "display_name": "Python 3.7.13 ('histo')",
   "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.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}