PyTorch vs TensorFlow Eager

pytorch_vs_tfe.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensorflow version: 1.7.0\n",
      "PyTorch version: 0.3.1.post2\n",
      "GPU: Tesla K80\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.eager.backprop import implicit_val_and_grad\n",
    "tf.enable_eager_execution()\n",
    "\n",
    "import torch\n",
    "from torch.autograd import Variable\n",
    "import torchvision\n",
    "\n",
    "print(f'Tensorflow version: {tf.__version__}')\n",
    "print(f'PyTorch version: {torch.__version__}')\n",
    "print(f'GPU: {torch.cuda.get_device_name(0)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Common"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "image_size = 224\n",
    "\n",
    "np_data = np.zeros((batch_size, 3, image_size, image_size), dtype=np.float32)\n",
    "np_target = np.zeros((batch_size), dtype=np.int32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_data = Variable(torch.from_numpy(np_data)).cuda()\n",
    "pt_target = Variable(torch.from_numpy(np_target)).cuda().long()\n",
    "pt_resnet = torchvision.models.resnet50().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "optim = torch.optim.SGD(pt_resnet.parameters(), lr=1e-3)\n",
    "criterion = torch.nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time per batch 0.67817147731781\n"
     ]
    }
   ],
   "source": [
    "torch.cuda.synchronize()\n",
    "start = time.time()\n",
    "for _ in range(100):\n",
    "    preds = pt_resnet(pt_data)\n",
    "    loss = criterion(preds, pt_target)\n",
    "    optim.zero_grad()\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "print(f'Time per batch {(time.time()  - start) / 100}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = tf.train.GradientDescentOptimizer(1e-3)\n",
    "def loss_fn(input, target):\n",
    "    return tf.losses.sparse_softmax_cross_entropy(target, tf_resnet(input))\n",
    "grad_fn = implicit_val_and_grad(loss_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time per batch 0.700717933177948\n"
     ]
    }
   ],
   "source": [
    "with tf.device('/gpu:0'):\n",
    "    tf_data = tf.constant(np_data.transpose(0, 2, 3, 1))\n",
    "    tf_resnet = tf.keras.applications.ResNet50(weights=None)\n",
    "    tf_target = tf.constant(np_target, dtype=tf.int32)\n",
    "    start = time.time()\n",
    "    for _ in range(100):\n",
    "        loss, grads = grad_fn(tf_data, tf_target)\n",
    "        optimizer.apply_gradients(grads)\n",
    "    print(f'Time per batch {(time.time()  - start) / 100}')"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}