argmax with CUDA in cupy vs pytorch vs tensorflow

CUDA Argmax Comparison.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "256x256 in batches of 128:\n",
    "- cp 112ms\n",
    "- th 570ms\n",
    "- tf 2332ms\n",
    "\n",
    "256x256 in batches of 256:\n",
    "- cp 96ms\n",
    "- th 307ms\n",
    "- tf 1438ms\n",
    "\n",
    "256x256 in batches of 1024:\n",
    "- cp 87ms\n",
    "- th 152ms\n",
    "- tf 2679ms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_rows = 256*256\n",
    "batch_size = 128\n",
    "cols = 256*256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "28609\n",
      "100.27ms 0.20ms/loop for 512 loops\n"
     ]
    }
   ],
   "source": [
    "import cupy as cp\n",
    "\n",
    "src_cp = cp.random.random((batch_size, cols)).astype(cp.float32)\n",
    "start = time()\n",
    "loops = total_rows//batch_size\n",
    "for i in range(loops):\n",
    "    out_cp = src_cp.argmax(axis=1)\n",
    "print(out_cp[0])\n",
    "duration = time() - start\n",
    "time_per_loop = duration / loops\n",
    "print(f'{1000*duration:.2f}ms {1000*time_per_loop:.2f}ms/loop for {loops} loops')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(64546, device='cuda:0')\n",
      "570.60ms 1.11ms/loop for 512 loops\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "src_th = torch.rand((batch_size, cols), device='cuda')\n",
    "start = time()\n",
    "loops = total_rows//batch_size\n",
    "for i in range(loops):\n",
    "    out_th = src_th.argmax(dim=1)\n",
    "print(out_th[0])\n",
    "duration = time() - start\n",
    "time_per_loop = duration / loops\n",
    "print(f'{1000*duration:.2f}ms {1000*time_per_loop:.2f}ms/loop for {loops} loops')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(8202, shape=(), dtype=int64)\n",
      "2330.52ms 4.55ms/loop for 512 loops\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "tf.enable_eager_execution()\n",
    "\n",
    "x_tf = tf.random_uniform((batch_size,cols))\n",
    "start = time()\n",
    "loops = total_rows//batch_size\n",
    "for i in range(loops):\n",
    "    out_tf = tf.argmax(x_tf, axis=1)\n",
    "print(out_tf[0])\n",
    "duration = time() - start\n",
    "time_per_loop = duration / loops\n",
    "print(f'{1000*duration:.2f}ms {1000*time_per_loop:.2f}ms/loop for {loops} loops')"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}