pytorch/parallelism_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nMulti-GPU examples\n==================\n\nData Parallelism is when we split the mini-batch of samples into\nmultiple smaller mini-batches and run the computation for each of the\nsmaller mini-batches in parallel.\n\nData Parallelism is implemented using ``torch.nn.DataParallel``.\nOne can wrap a Module in ``DataParallel`` and it will be parallelized\nover multiple GPUs in the batch dimension.\n\nDataParallel\n-------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch.nn as nn\n\n\nclass DataParallelModel(nn.Module):\n\n    def __init__(self):\n        super().__init__()\n        self.block1 = nn.Linear(10, 20)\n\n        # wrap block2 in DataParallel\n        self.block2 = nn.Linear(20, 20)\n        self.block2 = nn.DataParallel(self.block2)\n\n        self.block3 = nn.Linear(20, 20)\n\n    def forward(self, x):\n        x = self.block1(x)\n        x = self.block2(x)\n        x = self.block3(x)\n        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The code does not need to be changed in CPU-mode.\n\nThe documentation for DataParallel is\n`here <http://pytorch.org/docs/nn.html#torch.nn.DataParallel>`_.\n\n**Primitives on which DataParallel is implemented upon:**\n\n\nIn general, pytorch\u2019s `nn.parallel` primitives can be used independently.\nWe have implemented simple MPI-like primitives:\n\n- replicate: replicate a Module on multiple devices\n- scatter: distribute the input in the first-dimension\n- gather: gather and concatenate the input in the first-dimension\n- parallel\\_apply: apply a set of already-distributed inputs to a set of\n  already-distributed models.\n\nTo give a better clarity, here function ``data_parallel`` composed using\nthese collectives\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def data_parallel(module, input, device_ids, output_device=None):\n    if not device_ids:\n        return module(input)\n\n    if output_device is None:\n        output_device = device_ids[0]\n\n    replicas = nn.parallel.replicate(module, device_ids)\n    inputs = nn.parallel.scatter(input, device_ids)\n    replicas = replicas[:len(inputs)]\n    outputs = nn.parallel.parallel_apply(replicas, inputs)\n    return nn.parallel.gather(outputs, output_device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Part of the model on CPU and part on the GPU\n--------------------------------------------\n\nLet\u2019s look at a small example of implementing a network where part of it\nis on the CPU and part on the GPU\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class DistributedModel(nn.Module):\n\n    def __init__(self):\n        super().__init__(\n            embedding=nn.Embedding(1000, 10),\n            rnn=nn.Linear(10, 10).cuda(0),\n        )\n\n    def forward(self, x):\n        # Compute embedding on CPU\n        x = self.embedding(x)\n\n        # Transfer to GPU\n        x = x.cuda(0)\n\n        # Compute RNN on GPU\n        x = self.rnn(x)\n        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This was a small introduction to PyTorch for former Torch users.\nThere\u2019s a lot more to learn.\n\nLook at our more comprehensive introductory tutorial which introduces\nthe ``optim`` package, data loaders etc.: :doc:`/beginner/deep_learning_60min_blitz`.\n\nAlso look at\n\n-  :doc:`Train neural nets to play video games </intermediate/reinforcement_q_learning>`\n-  `Train a state-of-the-art ResNet network on imagenet`_\n-  `Train an face generator using Generative Adversarial Networks`_\n-  `Train a word-level language model using Recurrent LSTM networks`_\n-  `More examples`_\n-  `More tutorials`_\n-  `Discuss PyTorch on the Forums`_\n-  `Chat with other users on Slack`_\n\n\n"
      ]
    }
  ],
  "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.1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}