acceleration/multi_gpu_test.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi GPU Test\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Project-MONAI/tutorials/blob/main/acceleration/multi_gpu_test.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "!python -c \"import monai\" || pip install -q \"monai-weekly[ignite]\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MONAI version: 0.4.0\n",
      "Numpy version: 1.19.1\n",
      "Pytorch version: 1.7.0a0+7036e91\n",
      "MONAI flags: HAS_EXT = False, USE_COMPILED = False\n",
      "MONAI rev id: 0563a4467fa602feca92d91c7f47261868d171a1\n",
      "\n",
      "Optional dependencies:\n",
      "Pytorch Ignite version: 0.4.2\n",
      "Nibabel version: 3.2.1\n",
      "scikit-image version: 0.15.0\n",
      "Pillow version: 8.0.1\n",
      "Tensorboard version: 2.2.0\n",
      "gdown version: 3.12.2\n",
      "TorchVision version: 0.8.0a0\n",
      "ITK version: 5.1.2\n",
      "tqdm version: 4.54.1\n",
      "lmdb version: 1.0.0\n",
      "psutil version: 5.7.2\n",
      "\n",
      "For details about installing the optional dependencies, please visit:\n",
      "    https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Copyright 2020 MONAI Consortium\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "\n",
    "import torch\n",
    "from monai.config import print_config\n",
    "from monai.engines import create_multigpu_supervised_trainer\n",
    "from monai.networks.nets import UNet\n",
    "\n",
    "print_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_epochs = 2\n",
    "lr = 1e-3\n",
    "device = torch.device(\"cuda:0\")\n",
    "net = UNet(\n",
    "    spatial_dims=2,\n",
    "    in_channels=1,\n",
    "    out_channels=1,\n",
    "    channels=(16, 32, 64, 128, 256),\n",
    "    strides=(2, 2, 2, 2),\n",
    "    num_res_units=2,\n",
    ").to(device)\n",
    "\n",
    "\n",
    "def fake_loss(y_pred, y):\n",
    "    return (y_pred[0] + y).sum()\n",
    "\n",
    "\n",
    "def fake_data_stream():\n",
    "    while True:\n",
    "        yield torch.rand((10, 1, 64, 64)), torch.rand((10, 1, 64, 64))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1 GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "State:\n",
       "\titeration: 4\n",
       "\tepoch: 2\n",
       "\tepoch_length: 2\n",
       "\tmax_epochs: 2\n",
       "\toutput: 34253.46875\n",
       "\tbatch: <class 'tuple'>\n",
       "\tmetrics: <class 'dict'>\n",
       "\tdataloader: <class 'generator'>\n",
       "\tseed: <class 'NoneType'>\n",
       "\ttimes: <class 'dict'>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "opt = torch.optim.Adam(net.parameters(), lr)\n",
    "trainer = create_multigpu_supervised_trainer(net, opt, fake_loss, [device])\n",
    "trainer.run(fake_data_stream(), max_epochs=max_epochs, epoch_length=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### all GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "State:\n",
       "\titeration: 4\n",
       "\tepoch: 2\n",
       "\tepoch_length: 2\n",
       "\tmax_epochs: 2\n",
       "\toutput: 30694.720703125\n",
       "\tbatch: <class 'tuple'>\n",
       "\tmetrics: <class 'dict'>\n",
       "\tdataloader: <class 'generator'>\n",
       "\tseed: <class 'NoneType'>\n",
       "\ttimes: <class 'dict'>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "opt = torch.optim.Adam(net.parameters(), lr)\n",
    "trainer = create_multigpu_supervised_trainer(net, opt, fake_loss, None)\n",
    "trainer.run(fake_data_stream(), max_epochs=max_epochs, epoch_length=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "State:\n",
       "\titeration: 4\n",
       "\tepoch: 2\n",
       "\tepoch_length: 2\n",
       "\tmax_epochs: 2\n",
       "\toutput: 26988.939453125\n",
       "\tbatch: <class 'tuple'>\n",
       "\tmetrics: <class 'dict'>\n",
       "\tdataloader: <class 'generator'>\n",
       "\tseed: <class 'NoneType'>\n",
       "\ttimes: <class 'dict'>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = net.to(torch.device(\"cpu:0\"))\n",
    "opt = torch.optim.Adam(net.parameters(), lr)\n",
    "trainer = create_multigpu_supervised_trainer(net, opt, fake_loss, [])\n",
    "trainer.run(fake_data_stream(), max_epochs=max_epochs, epoch_length=2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}