fft_keras_vs_pytorch_gpu.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Benchmarking keras vs pytorch on IFFT2D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I am currently working on MRI reconstruction deep learning models. A lot of these models involve a heavy use of FFT2D and IFFT2D allowing to correct the artifacts both in image space and k-space alternatively.\n",
    "\n",
    "Therefore, using a deep learning framework which allows for an efficient and easy-to-use implementation of the IFFT2D/FFT2D pair. I was using `keras` with a `tensorflow` backend originally because I was familiar with it and because I thought there was no reason for the Fourier transform operations to be inefficient.\n",
    "\n",
    "It turns out that the Fourier transform (at least in 2D) is particularly inefficient in `tensorflow` on CPU. You can follow this [Github issue](https://github.com/tensorflow/tensorflow/issues/6541) to learn more (it's closed, but the discussion is actually going on).\n",
    "\n",
    "I thus decided to see whether `pytorch` would perform better. The answer, according to this simple benchmark, is that it is almost 4 times faster than `keras` with a `tensorflow` backend.\n",
    "\n",
    "This benchmark is to be run on GPU (but can easily be tweaked to be run on CPU, where tensorflow is 40 times slower). You will obtain trace timelines for the prediction part for both `pytorch` and `keras`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # this just to make sure we are using only on CPU\n",
    "# import os\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"-1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this to get rid of tensorflow deprecation warnings\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import logging\n",
    "logging.getLogger('tensorflow').disabled = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "from keras.layers import Input, Lambda, Conv2D, concatenate\n",
    "from keras.models import Model\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.signal import ifft2d\n",
    "from tensorflow.python.client import timeline\n",
    "import torch\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.3.0.dev20190828'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.14.0'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.2.4'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import keras\n",
    "keras.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fake data creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = torch.float32\n",
    "device = torch.device(\"cuda\")\n",
    "\n",
    "N, im_size = 35, 320\n",
    "# Create random input and output data\n",
    "x = torch.randn(N, im_size, im_size, 2, device=device, dtype=dtype)\n",
    "y = torch.randn(N, im_size, im_size, device=device, dtype=dtype)\n",
    "\n",
    "# in numpy\n",
    "x_np = x.cpu().numpy()\n",
    "x_np = x_np[...,0] + 1j * x_np[..., 1]\n",
    "x_np = x_np[..., None]\n",
    "y_np = y.cpu().numpy()\n",
    "y_np = y_np[..., None]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model definition and creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvInvNet(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ConvInvNet, self).__init__()\n",
    "        self.conv = nn.Conv2d(2, 2, kernel_size=3, padding=1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        y = torch.ifft(x, 2)\n",
    "        # this because pytorch doesn't support NHWC\n",
    "        y = y.permute(0, 3, 1, 2)\n",
    "        y = self.conv(y)\n",
    "        y = y.permute(0, 2, 3, 1)\n",
    "        y = (y ** 2).sum(dim=-1).sqrt()\n",
    "        return y\n",
    "    \n",
    "model_conv = ConvInvNet().cuda()\n",
    "criterion = torch.nn.MSELoss(reduction='sum')\n",
    "optimizer = torch.optim.SGD(model_conv.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.86 ms, sys: 10.8 ms, total: 16.7 ms\n",
      "Wall time: 34.2 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# predicting\n",
    "r = model_conv(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8.45 ms, sys: 5.99 ms, total: 14.4 ms\n",
      "Wall time: 13 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "with torch.autograd.profiler.profile(use_cuda=True) as prof:\n",
    "    r = model_conv(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = f'fft_pytorch_timeline_bs{N}_predict.json'\n",
    "prof.export_chrome_trace(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 78.9 ms, sys: 19.4 ms, total: 98.3 ms\n",
      "Wall time: 108 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# training\n",
    "for _ in range(5):\n",
    "    r = model_conv(x)\n",
    "    loss = criterion(r, np.squeeze(y))\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Keras"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model definition and creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def concatenate_real_imag(x):\n",
    "    x_real = Lambda(tf.math.real)(x)\n",
    "    x_imag = Lambda(tf.math.imag)(x)\n",
    "    return concatenate([x_real, x_imag])\n",
    "\n",
    "def to_complex(x):\n",
    "    return tf.complex(x[0], x[1])\n",
    "\n",
    "def complex_from_half(x, n, output_shape):\n",
    "    return Lambda(lambda x: to_complex([x[..., :n], x[..., n:]]), output_shape=output_shape)(x)\n",
    "\n",
    "input_size = (320, None, 1)\n",
    "kspace_input = Input(input_size, dtype='complex64', name='kspace_input')\n",
    "inv_kspace = Lambda(ifft2d, output_shape=input_size)(kspace_input)\n",
    "inv_kspace = concatenate_real_imag(inv_kspace)\n",
    "inv_kspace = Conv2D(\n",
    "    2,\n",
    "    3,\n",
    "    activation='linear',\n",
    "    padding='same',\n",
    "    kernel_initializer='he_normal',\n",
    ")(inv_kspace)\n",
    "inv_kspace = complex_from_half(inv_kspace, 1, input_size)\n",
    "abs_inv_kspace = Lambda(tf.math.abs)(inv_kspace)\n",
    "model_conv_keras = Model(inputs=kspace_input, outputs=abs_inv_kspace)\n",
    "run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)\n",
    "run_metadata = tf.RunMetadata()\n",
    "\n",
    "model_conv_keras.compile(\n",
    "    optimizer='adam',\n",
    "    loss='mse',\n",
    "    options=run_options,\n",
    "    run_metadata=run_metadata,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.85 s, sys: 1.36 s, total: 3.21 s\n",
      "Wall time: 3.23 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# predicting\n",
    "r = model_conv_keras.predict_on_batch(x_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 33 ms, sys: 20.6 ms, total: 53.6 ms\n",
      "Wall time: 48.6 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# predicting\n",
    "r = model_conv_keras.predict_on_batch(x_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# profiling output\n",
    "tl = timeline.Timeline(run_metadata.step_stats)\n",
    "ctf = tl.generate_chrome_trace_format()\n",
    "with open(f'fft_keras_timeline_bs{N}_predict.json', 'w') as f:\n",
    "    f.write(ctf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 670 ms, sys: 77.9 ms, total: 748 ms\n",
      "Wall time: 714 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# training\n",
    "r = model_conv_keras.fit(x_np, y_np, batch_size=35, epochs=5, verbose=0)"
   ]
  },
  {
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}