train_cifar.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d2c01019-b574-448e-98a4-a7cfe3105a53",
   "metadata": {
    "tags": []
   },
   "source": [
    "# CIFAR-10 with Ray Train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a254aa55-db54-4cea-a9c6-e125ff5ff038",
   "metadata": {},
   "source": [
    "This demo walks you through the lifecycle of deep learning training job from beginning to end.\n",
    "\n",
    "At a high level, we will go through 3 steps:\n",
    "1. Run your training script as a non-distributed script, locally.\n",
    "2. Convert the script to run with Ray Train in a distributed manner, again locally.\n",
    "3. Run the same script in a distributed manner, on a cluster with GPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0b21fdf-5597-40c0-944a-59d8d305f51e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83dbfc1d-ca2c-497c-800a-9c50313b7dd1",
   "metadata": {},
   "source": [
    "## Local Non-Distributed Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d04b19b3-8b16-4a5e-b0f2-1876007b11dc",
   "metadata": {},
   "source": [
    "Let's start by running our training function. Note that we pass in two configuration values:\n",
    "1. `\"test_mode\": True`: This cuts down the training and test datasets to 64 records (from 5000 and 1000).\n",
    "2. `\"num_epochs\": 10`: We will train for just 10 epochs to verify it's running.\n",
    "\n",
    "While it's running, we can dive into the code a bit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68ea8e3f-83e7-4cf5-88de-dde9e12adcb4",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from utils import train_func\n",
    "import time\n",
    "\n",
    "start_time=time.time()\n",
    "results = train_func(config={\"test_mode\": True, \"num_epochs\": 10})\n",
    "print(f\"--- Finished training in {time.time() - start_time} seconds ---\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fffc215-02dc-4201-bc43-07f0af6741d4",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63511213-c569-4e50-bcc9-85c23800cd0f",
   "metadata": {},
   "source": [
    "Now that we've \"trained\" a model, let's test it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45721705-c4c3-4786-af93-ee6aa7c5778d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import get_model, get_test_data, predict_and_display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c861f7c-447d-4064-8563-24613cbe9720",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_state_dict = results[\"model_state_dict\"]\n",
    "model = get_model(model_state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5588c61e-20db-4f3d-9d49-1d9e3fc9f6f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_images, test_labels = get_test_data(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1df51680-d1ac-4f18-babd-da1731631e6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_and_display(model, test_images, test_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39826476-e715-46bb-8c80-4a73f2036238",
   "metadata": {},
   "source": [
    "## Local Distributed Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c54ac4fb-201a-45ec-ac85-b04f35a00f84",
   "metadata": {},
   "source": [
    "Now that we have a working training function, our next step is to distribute it. This is where we'll start using Ray Train.\n",
    "\n",
    "To move to Ray Train, let's revisit the python script and uncomment out the Ray Train utility functions."
   ]
  },
  {
   "cell_type": "raw",
   "id": "fca0d153-f17d-4842-b2c8-8d1f22c5954d",
   "metadata": {},
   "source": [
    "%load code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6275aa22-c3dd-48af-b358-fc37b48eddcd",
   "metadata": {},
   "source": [
    "Now that our training function is set up, we'll now need to set up a Ray Train `Trainer`, the main class you will interface with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22115fab-9df0-45b3-8079-5279d1365620",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.train import Trainer\n",
    "trainer = Trainer(backend=\"torch\", num_workers=2, logdir=\"~/workspace/scratch/cifar10/results/local\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d846dcff-8f92-47be-bc1a-f52088fb0bb3",
   "metadata": {},
   "source": [
    "A quick summary of the arguments:\n",
    "1. `backend=\"torch\"` - Distributed PyTorch environment will be set up on the worker processes.\n",
    "2. `num_workers=2` - Training will be split across two worker processes.\n",
    "3. (optional) `logdir=\"~/workspace/scratch/cifar10/results/local\"` - Setup logs locally for easier viewing."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e63633a4-e56b-4d67-b6f8-8acacdd21ec4",
   "metadata": {},
   "source": [
    "After setting up the `Trainer`, we call `start` to start the worker processes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b43cdea-04b6-495c-ae05-1bab88b1a680",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a55c4af-c4a6-4e28-a0e0-f76fa4598f2a",
   "metadata": {},
   "source": [
    "Now that we have 2 workers up, let's run the updated training function!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1efd411-59e9-4c46-b3d1-626885161b61",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "start_time=time.time()\n",
    "trainer.run(train_func, config={\"test_mode\": True, \"num_epochs\": 10})\n",
    "print(f\"--- Finished training in {time.time() - start_time} seconds ---\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f47e39f1-345a-43c2-9eb4-56eb387fc4f8",
   "metadata": {},
   "source": [
    "Let's take a look at some of the attributes that can be accessed after training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6954863b-f45c-4493-a8ef-57d86c4a3cc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.latest_run_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63f6a66e-f2fd-4880-b270-4b02b3715645",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.latest_checkpoint_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6351b72c-33c2-4373-9504-5ffd71c0f380",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer.latest_checkpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "896fdb60-3b18-4e11-b345-c68b12d3a607",
   "metadata": {},
   "source": [
    "Now let's try _this_ \"trained\" model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35307e7e-a536-45de-9244-44cce0a99303",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = trainer.latest_checkpoint\n",
    "model_state_dict = checkpoint[\"model_state_dict\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbe4a6e7-7f1f-4c99-9dd4-a2ce34d563c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = get_model(model_state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb9a23c-bcb8-4d6d-b0cf-f603bc481f3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_and_display(model, test_images, test_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07cd6904-92a8-4040-8396-6bece0401977",
   "metadata": {},
   "source": [
    "To end this step, we'll clean up the worker processes and resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "684310fa-63ad-4fea-b899-ad12d2f5d246",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.shutdown()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82245a74-33ac-4039-8180-c567e714bae2",
   "metadata": {},
   "source": [
    "## GPU Cluster Distributed Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "158cb8f8-699d-4520-a3cc-d825685ef66b",
   "metadata": {},
   "source": [
    "At this point, we've verified that our code runs locally (and distributedly!) for a small subset of the data. But what we _really_ want is to scale up to a cluster of GPUs so that we can train on the entire dataset.\n",
    "\n",
    "Let's setup that Ray Cluster!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28ab151e-cd2b-416e-83a5-13f56c189cd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load cluster-config.yaml\n",
    "cluster_name: cifar10\n",
    "\n",
    "max_workers: 8\n",
    "\n",
    "provider:\n",
    "    type: aws\n",
    "    region: us-west-1\n",
    "\n",
    "available_node_types:\n",
    "    cpu-node:\n",
    "        min_workers: 0\n",
    "        max_workers: 0\n",
    "        node_config:\n",
    "            InstanceType: m5.4xlarge\n",
    "        resources: {}\n",
    "    1-gpu-node:\n",
    "        min_workers: 8\n",
    "        max_workers: 8\n",
    "        node_config:\n",
    "            InstanceType: g4dn.xlarge\n",
    "            InstanceMarketOptions:\n",
    "                MarketType: spot\n",
    "        resources: {}\n",
    "\n",
    "head_node_type: cpu-node\n",
    "\n",
    "setup_commands:\n",
    "    - pip install ray[default] torch torchvision\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "303de909-7f1f-4333-9bc9-ad762bddfc4b",
   "metadata": {},
   "source": [
    "In a separate terminal:\n",
    "```bash\n",
    "# Launch your cluster\n",
    "$ ray up cluster-config.yaml\n",
    "\n",
    "# SSH into the cluster and enable port forwarding\n",
    "$ ray attach cluster-config.yaml -p 10001\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "661422a0-e261-4d62-9031-23b24b1a4953",
   "metadata": {},
   "source": [
    "At this point, you're ready to run your script on the cluster!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bc6a17f-457c-435e-97d5-a318e9914695",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray\n",
    "ray_address=\"ray://localhost:10001\"\n",
    "# ray_address=\"anyscale://my-cluster\"\n",
    "ray.init(address=ray_address, runtime_env={\"working_dir\": \".\", \"excludes\":[\"data\", \"results\"]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a07a2a90-a572-4116-bad2-85b441a06306",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(backend=\"torch\", num_workers=8, use_gpu=True, logdir=\"~/workspace/scratch/cifar10/results/cluster\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2669f5bf-a901-4dc9-8ea5-8eb01b15dabe",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "560643ee-264a-41c2-8bff-0fc0e6883999",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ray.train.callbacks import JsonLoggerCallback, TBXLoggerCallback\n",
    "\n",
    "# Scale it up!\n",
    "config = {\"batch_size\": 1024, \"lr\": 0.1, \"num_epochs\": 100}\n",
    "callbacks = [JsonLoggerCallback(), TBXLoggerCallback()]\n",
    "\n",
    "trainer.run(train_func, config=config, callbacks=callbacks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7e10ef9-cf57-4423-b79a-9e1466762fd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "! tensorboard dev upload --one_shot --logdir /Users/matt/workspace/scratch/cifar10/results/cluster/run_001 --name cifar10-demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e0c579-d882-41bc-a1ef-c1edff60a41e",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer.latest_checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49f8fad8-6d27-46e9-aec8-a2d2fba8f191",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = trainer.latest_checkpoint\n",
    "model_state_dict = checkpoint[\"model_state_dict\"]\n",
    "model = get_model(model_state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "515357d1-991a-4133-89ad-533cd4bf97ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_and_display(model, test_images, test_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "881c66c8-8301-4db8-a0f0-3cf4e2fc9902",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.shutdown()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cifar10",
   "language": "python",
   "name": "cifar10"
  },
  "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}