Demonstrating Convolution-based Upsampling.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "\n",
    "import upsampling.layers as L"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validating that custom convolution-based upsampling layers are equivalent to PyTorch implementations\n",
    "\n",
    "These layers are written for code clarity, not for speed. Each layer is implemented using nested for-loops, so it is best to use small values for feature map dimensions, upsampling factors, and kernel sizes. Note that the sub-pixel and resize convolutions use same-paded convolution, so the padding is a factor of the kernel size $K$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Conv2            pixel-to-pixel = 100.0%\n",
      "PixelShuffle     pixel-to-pixel = 100.0%\n",
      "Sub-Pixel Conv   pixel-to-pixel = 100.0%\n",
      "NN Resize Conv   pixel-to-pixel = 100.0%\n",
      "Deconvolution    pixel-to-pixel = 100.0%\n"
     ]
    }
   ],
   "source": [
    "H = W = 4 # height/width - note using square feature maps\n",
    "C = 3     # number of channels - note that the squeeze() operation below means this check only works for C > 1\n",
    "K = 3     # kernel_size\n",
    "N = 1     # batch_size\n",
    "r = 2     # upscaling factor\n",
    "P = (K - 1) // 2 # <---- needs to be same-padded to be valid\n",
    "\n",
    "def p2p_accuracy(source:torch.Tensor, target:torch.Tensor, n_digits:int = 5) -> float:\n",
    "    distance = L.round(source - target, n_digits).abs()\n",
    "    correct = (distance == 0).float()\n",
    "    return correct.mean()\n",
    "\n",
    "# ------------------------------------------------------------------------------------- #\n",
    "# Testing the Conv2d Operator\n",
    "x = torch.randn(N, C, H, W)\n",
    "\n",
    "m = nn.Conv2d(in_channels=C, out_channels=C, kernel_size=K, bias=False, padding=P)\n",
    "n = L.Conv2d(in_channels=C, out_channels=C, kernel_size=K, padding=P)\n",
    "n.weight = m.weight\n",
    "acc = p2p_accuracy(m(x).squeeze(), n(x.squeeze(0)))\n",
    "print(f\"Conv2            pixel-to-pixel = {acc:.1%}\")\n",
    "\n",
    "# ------------------------------------------------------------------------------------- #\n",
    "# Testing the PixelShuffle Operator\n",
    "x = torch.zeros(N, C*(r**2), H, W)\n",
    "\n",
    "for c in range(C*(r**2)):\n",
    "    x[:,c,:,:] = c\n",
    "\n",
    "m = nn.PixelShuffle(r)\n",
    "n = L.PixelShuffle(r)\n",
    "acc = p2p_accuracy(m(x).squeeze(), n(x.squeeze(0)))\n",
    "print(f\"PixelShuffle     pixel-to-pixel = {acc:.1%}\")\n",
    "\n",
    "# ------------------------------------------------------------------------------------- #\n",
    "# Testing the Sub-Pixel Convolution Operator\n",
    "x = torch.randn(N, C, H, W)\n",
    "m_layer1 = nn.Conv2d(in_channels=C, out_channels=C*(r**2), kernel_size=K, padding=P, bias=False)\n",
    "m_layer2 = nn.PixelShuffle(r)\n",
    "p = m_layer2(m_layer1(x)).squeeze()\n",
    "\n",
    "n = L.SubPixelConvolution(in_channels=C, out_channels=C*(r**2), kernel_size=K, scaling_factor=r, padding=P)\n",
    "n.convolution.weight = m_layer1.weight\n",
    "q = n(x.squeeze(0))\n",
    "acc = p2p_accuracy(p, q)\n",
    "\n",
    "print(f\"Sub-Pixel Conv   pixel-to-pixel = {acc:.1%}\")\n",
    "\n",
    "x = torch.randn(N, C, H, W)\n",
    "m = nn.Upsample(scale_factor=r, mode='nearest')\n",
    "n = nn.Conv2d(in_channels=C, out_channels=C, kernel_size=K, padding=P, stride=1, bias=False)\n",
    "p = n(m(x))\n",
    "\n",
    "resize_conv = L.ResizeConvolution(scaling_factor=r, in_channels=C, out_channels=C, kernel_size=K, padding=P, stride=1)\n",
    "resize_conv.convolution.weight = n.weight\n",
    "q = resize_conv(x.squeeze(0))\n",
    "acc = p2p_accuracy(p, q)\n",
    "\n",
    "print(f\"NN Resize Conv   pixel-to-pixel = {acc:.1%}\")\n",
    "\n",
    "# ------------------------------------------------------------------------------------- #\n",
    "# Testing the Deconvolution Operator\n",
    "x = torch.randn(N, C, H, W)\n",
    "m = nn.ConvTranspose2d(in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=P*r, bias=False)\n",
    "n = L.Deconvolution(in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=P*r)\n",
    "n.weight = m.weight\n",
    "acc = p2p_accuracy(m(x).squeeze(), n(x.squeeze(0)))\n",
    "\n",
    "print(f\"Deconvolution    pixel-to-pixel = {acc:.1%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Showing that a sub-pixel convolution is equivalent to a deconvolution using the weight shuffle algorithm\n",
    "\n",
    "The learned kernels of the sub-pixel convolution can be transformed into those of a functionally equivalent deconvolution using `WeightShuffle`.\n",
    "\n",
    "**More on Sub-pixel Convolution**\n",
    "- [Shi *et al.* (2016)- Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network](https://arxiv.org/abs/1609.05158)\n",
    "\n",
    "- [Shi *et al.* (2016) - Is the deconvolution layer the same as a convolutional layer?](https://arxiv.org/abs/1609.07009)\n",
    "\n",
    "- [Aitken *et al.* (2017) - Checkerboard artifact free sub-pixel convolution: A note on sub-pixel convolution, resize convolution and convolution resize](https://arxiv.org/abs/1707.02937)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sub-pixel Convolution versus Deconvolution    pixel-to-pixel = 100.0%\n"
     ]
    }
   ],
   "source": [
    "H = W = 2\n",
    "K = 3\n",
    "r = 2\n",
    "C = 1\n",
    "P = (K - 1) // 2 # <---- needs to be same-padded to be valid\n",
    "\n",
    "x = torch.randn(N, C, H, W)\n",
    "\n",
    "# Sub-pixel convolution operation\n",
    "subpixel_conv = L.SubPixelConvolution(in_channels=C, out_channels=C*(r**2), kernel_size=K, scaling_factor=r, padding=P)\n",
    "\n",
    "# Deconvolution operation\n",
    "deconvolution = L.Deconvolution(in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=P*r)\n",
    "weight_shuff  = L.WeightShuffle(r)\n",
    "\n",
    "# Shuffle convolution weights to be equivalent to the deconvolution\n",
    "deconvolution.weight = weight_shuff(subpixel_conv.convolution.weight)\n",
    "\n",
    "# Run and compare\n",
    "y_conv   = subpixel_conv(x.squeeze(0))\n",
    "y_deconv = deconvolution(x.squeeze(0))\n",
    "\n",
    "acc = p2p_accuracy(y_conv, y_deconv)\n",
    "\n",
    "print(f\"Sub-pixel Convolution versus Deconvolution    pixel-to-pixel = {acc:.1%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Showing that a resize convolution is equivalent to a deconvolution using the weight convolution algorithm\n",
    "\n",
    "The learned kernels of a nearest neighbor resize convolution can also be transformed into that of a functionally equivalent deconvolution using `WeightConvolution`.\n",
    "\n",
    "\n",
    "**More on Resize Convolution**\n",
    "- [Dong *et al* (2015) - Image Super-Resolution Using Deep Convolutional Networks](https://ieeexplore.ieee.org/abstract/document/7115171) - using bicubic interpolation\n",
    "- [Odena *et al* (2016) - Deconvolution and Checkerboard Artifacts](https://distill.pub/2016/deconv-checkerboard/) - using nearest neighbor or bilinear interpolation\n",
    "\n",
    "- [Aitken *et al.* (2017) - Checkerboard artifact free sub-pixel convolution](https://arxiv.org/abs/1707.02937) - using nearest neighbor interpolation\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NN resize convolution versus deconvolution    pixel-to-pixel = 100.0%\n"
     ]
    }
   ],
   "source": [
    "H = W = 2\n",
    "K = 3\n",
    "r = 2\n",
    "C = 1\n",
    "P = (K - 1) // 2 # <---- needs to be same-padded to be valid\n",
    "\n",
    "x = torch.randn(N, C, H, W)\n",
    "\n",
    "# NN resize convolution operation\n",
    "resize_conv = L.ResizeConvolution(scaling_factor=r, in_channels=C, out_channels=C, kernel_size=K, padding=P, stride=1)\n",
    "\n",
    "# Deconvolution operation\n",
    "deconvolution = L.Deconvolution(in_channels=C, out_channels=C, kernel_size=r + K - 1, stride=r, padding=P)\n",
    "\n",
    "# Convolved the convolution weights to be equivalent to the deconvolution\n",
    "deconvolution.weight = L.weight_convolution(resize_conv.convolution.weight,\n",
    "                                            in_channels=C,\n",
    "                                            out_channels=C,\n",
    "                                            kernel_size=K,\n",
    "                                            scaling_factor=r)\n",
    "\n",
    "# Run and compare\n",
    "y_conv   = resize_conv(x.squeeze(0))\n",
    "y_deconv = deconvolution(x.squeeze(0))\n",
    "\n",
    "acc = p2p_accuracy(y_conv, y_deconv)\n",
    "\n",
    "print(f\"NN resize convolution versus deconvolution    pixel-to-pixel = {acc:.1%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Showing that the low-level deconvolution algorithms give the identical results\n",
    "\n",
    "The standard deconvolution algorithm strides over the input space, creating overlapping sums in the output space when the stride $S$ is smaller than the kernel size $K$. In the context of training, these overlapping sums have been shown to introduce checkerboard artifacts as a result of gradient updates. In the context of inference, accumulating over these overlapping regions requires storing partial sums when $K > S$ which leads to significant overhead.  Algorithmic approaches to work around this *overlapping sums* problem are divided into three categories - reverse looping deconvolution (REVD), fractionally strided deconvolution (STRD), and transforming deconvolution to convolution (TDC).\n",
    "\n",
    "\n",
    "**More on standard deconvolution (STDD)**\n",
    "- [Dumoulin *et al.* (2016) - A guide to convolution arithmetic for deep learning](https://arxiv.org/abs/1603.07285)\n",
    "\n",
    "**More on reverse looping deconvolution (REVD)**\n",
    "\n",
    "- [Zhang *et al.* (2017) - A Design Methodology for Efficient Implementation of Deconvolutional Neural Networks on an FPGA](https://arxiv.org/abs/1705.02583)\n",
    "\n",
    "- [Colbert *et al.* (2021) - A Competitive Edge: Can FPGAs Beat GPUs at DCNN Inference Acceleration in Resource-Limited Edge Computing Applications?](https://arxiv.org/abs/2102.00294)\n",
    "\n",
    "**More on improved reverse looping deconvolution (REVD2)**\n",
    "- [Colbert *et al.* (2021) - An Energy-Efficient Edge Computing Paradigm for Convolution-based Image Upsampling](https://arxiv.org/abs/2107.07647)\n",
    "\n",
    "**More on transforming deconvolution to convolution (TDC)**\n",
    "- [Chang *et al.* (2018) - An Energy-Efficient FPGA-Based Deconvolutional Neural Networks Accelerator for Single Image Super-Resolution](https://ieeexplore.ieee.org/abstract/document/8584497)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pixel-to-pixel accuracy = 100.0%\n"
     ]
    }
   ],
   "source": [
    "# ------------------------------------------------------------------------------------- #\n",
    "# Testing the Standard Deconvolution algorithm against the Reverse Deconvolution algorithm\n",
    "x = torch.randn(N, C, H, W)\n",
    "m = L.Deconvolution(\n",
    "    in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=r, algorithm=L.DeconvolutionAlgorithms.STDD\n",
    ")\n",
    "n = L.Deconvolution(\n",
    "    in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=r, algorithm=L.DeconvolutionAlgorithms.REVD\n",
    ")\n",
    "n.weight = m.weight = m.weight\n",
    "\n",
    "acc = p2p_accuracy(n(x.squeeze(0)), m(x.squeeze(0)))\n",
    "\n",
    "print(f\"Pixel-to-pixel accuracy = {acc:.1%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pixel-to-pixel accuracy = 100.0%\n"
     ]
    }
   ],
   "source": [
    "# ------------------------------------------------------------------------------------- #\n",
    "# Testing the Standard Deconvolution algorithm against the Reverse Deconvolution-2 algorithm\n",
    "x = torch.randn(N, C, H, W)\n",
    "m = L.Deconvolution(\n",
    "    in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=r, algorithm=L.DeconvolutionAlgorithms.STDD\n",
    ")\n",
    "n = L.Deconvolution(\n",
    "    in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=r, algorithm=L.DeconvolutionAlgorithms.REVD2\n",
    ")\n",
    "n.weight = m.weight = m.weight\n",
    "\n",
    "acc = p2p_accuracy(n(x.squeeze(0)), m(x.squeeze(0)))\n",
    "\n",
    "print(f\"Pixel-to-pixel accuracy = {acc:.1%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pixel-to-pixel accuracy = 100.0%\n"
     ]
    }
   ],
   "source": [
    "# ------------------------------------------------------------------------------------- #\n",
    "# Testing the Standard Deconvolution algorithm against the Fractionally Strided Convolution algorithm\n",
    "x = torch.randn(N, C, H, W)\n",
    "m = L.Deconvolution(\n",
    "    in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=r, algorithm=L.DeconvolutionAlgorithms.STDD\n",
    ")\n",
    "n = L.Deconvolution(\n",
    "    in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=r, algorithm=L.DeconvolutionAlgorithms.STRD\n",
    ")\n",
    "n.weight = m.weight\n",
    "\n",
    "acc = p2p_accuracy(n(x.squeeze(0)), m(x.squeeze(0)))\n",
    "\n",
    "print(f\"Pixel-to-pixel accuracy = {acc:.1%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pixel-to-pixel accuracy = 100.0%\n"
     ]
    }
   ],
   "source": [
    "# ------------------------------------------------------------------------------------- #\n",
    "# Testing the Standard Deconvolution algorithm against Transforming Deconvolution to Convolution\n",
    "x = torch.randn(N, C, H, W)\n",
    "m = L.Deconvolution(\n",
    "    in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=r, algorithm=L.DeconvolutionAlgorithms.STDD\n",
    ")\n",
    "n = L.Deconvolution(\n",
    "    in_channels=C, out_channels=C, kernel_size=K*r, stride=r, padding=r, algorithm=L.DeconvolutionAlgorithms.TDC\n",
    ")\n",
    "n.weight = m.weight\n",
    "\n",
    "acc = p2p_accuracy(n(x.squeeze(0)), m(x.squeeze(0)))\n",
    "\n",
    "print(f\"Pixel-to-pixel accuracy = {acc:.1%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}