experimental/using-pytorch-lightning/1.train-single-node.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train with PyTorch Lightning\n",
    "\n",
    "description: train single-node, including single-node multi-gpu, pytorch lightning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade tensorboard azureml-tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from azureml.core import Workspace\n",
    "\n",
    "ws = Workspace.from_config()\n",
    "ws"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# training script\n",
    "source_dir = \"src\"\n",
    "script_name = \"train.py\"\n",
    "\n",
    "# environment file\n",
    "environment_file = \"environment.yml\"\n",
    "\n",
    "# azure ml settings\n",
    "environment_name = \"pt-lightning\"\n",
    "experiment_name = \"pt-lightning-tutorial\"\n",
    "compute_name = \"gpu-K80-2\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create environment\n",
    "\n",
    "Define a conda environment YAML file with your training script dependencies and create an Azure ML environment. The dependencies for this tutorial include **torch**, **torchvision**, and **pytorch-lightning**.\n",
    "\n",
    "Since this example is for GPU training, you will need to specify a GPU base image that has the necessary dependencies. Azure ML maintains a set of base images published on Microsoft Container Registry (MCR) that you can use, see the [Azure/AzureML-Containers](https://github.com/Azure/AzureML-Containers) GitHub repo for more information.\n",
    "\n",
    "Azure ML will build a conda environment with the dependencies you specified in your .yml file on the base image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from azureml.core import Environment\n",
    "\n",
    "env = Environment.from_conda_specification(environment_name, environment_file)\n",
    "\n",
    "# specify a GPU base image\n",
    "env.docker.enabled = True\n",
    "env.docker.base_image = (\n",
    "    \"mcr.microsoft.com/azureml/openmpi3.1.2-cuda10.2-cudnn8-ubuntu18.04\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, you can just capture all your dependencies directly in a custom Docker image or Dockerfile, and create your environment from that. For more information, see [Train with custom image](https://docs.microsoft.com/azure/machine-learning/how-to-train-with-custom-image)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure and run training job\n",
    "Create a ScriptRunConfig to specify the training script & arguments, environment, and cluster to run on.\n",
    "\n",
    "For single-node, single-GPU training, specify `1` GPU to the `--gpus` command-line argument expected by Lightning.\n",
    "Note that you do not need to define this flag manually in your training script as Lightning can add it automatically. The training script parses the command-line arguments and passes them to the [`Trainer()`](https://pytorch-lightning.readthedocs.io/en/stable/trainer.html?highlight=Trainer).\n",
    "\n",
    "Lightning handles all the NVIDIA flags for you, there's no need to set them yourself. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from azureml.core import ScriptRunConfig, Experiment\n",
    "\n",
    "src = ScriptRunConfig(\n",
    "    source_directory=source_dir,\n",
    "    script=script_name,\n",
    "    arguments=[\"--max_epochs\", 25, \"--gpus\", 1],\n",
    "    compute_target=compute_name,\n",
    "    environment=env,\n",
    ")\n",
    "\n",
    "run = Experiment(ws, experiment_name).submit(src)\n",
    "run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run.wait_for_completion(show_output=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Single-node multi-GPU training\n",
    "\n",
    "Lightning supports several [distributed modes](https://pytorch-lightning.readthedocs.io/en/stable/multi_gpu.html#distributed-modes) for training. DistributedDataParallel (DDP) is recommended over DataParallel (DP) for training.\n",
    "\n",
    "For multi-GPU training on a single node, specify the number of GPUs to train on (typically this will correspond to the number of GPUs in your cluster's SKU) and the distributed mode, in this case DistributedDataParallel (\"ddp\"), which Lightning expects as arguments `--gpus` and `--accelerator`, respectively. The Lightning implementation of DDP will manage starting the individual processes on each GPU under the hood. See their [Multi-GPU](https://pytorch-lightning.readthedocs.io/en/stable/multi_gpu.html) training documentation for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from azureml.core import ScriptRunConfig, Experiment\n",
    "\n",
    "src = ScriptRunConfig(\n",
    "    source_directory=source_dir,\n",
    "    script=script_name,\n",
    "    arguments=[\"--max_epochs\", 25, \"--gpus\", 2, \"--accelerator\", \"ddp\"],\n",
    "    compute_target=compute_name,\n",
    "    environment=env,\n",
    ")\n",
    "\n",
    "run = Experiment(ws, experiment_name).submit(src)\n",
    "run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can monitor the progress of the run with a Jupyter widget. Like the run submission, the widget is asynchronous and provides live updates every 10-15 seconds until the job completes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from azureml.widgets import RunDetails\n",
    "\n",
    "RunDetails(run).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run.wait_for_completion(show_output=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8",
   "language": "python",
   "name": "python3.8"
  },
  "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.5-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}