examples/out/train-resnet50-multiray-wds.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "577bba5c",
   "metadata": {
    "papermill": {
     "duration": 0.002808,
     "end_time": "2023-12-18T08:08:20.493793",
     "exception": false,
     "start_time": "2023-12-18T08:08:20.490985",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# WebDataset + Distributed PyTorch Training\n",
    "\n",
    "This notebook illustrates how to use the Web Indexed Dataset (`wids`) library for distributed PyTorch training using `DistributedDataParallel`.\n",
    "\n",
    "Using `webdataset` results in training code that is almost identical to plain PyTorch except for the dataset creation.\n",
    "Since `WebDataset` is an iterable dataset, you need to account for that when creating the `DataLoader`. Furthermore, for\n",
    "distributed training, easy restarts, etc., it is convenient to use a resampled dataset; this is in contrast to\n",
    "sampling without replacement for each epoch as used more commonly for small, local training. (If you want to use\n",
    "sampling without replacement with webdataset format datasets, see the companion `wids`-based training notebooks.)\n",
    "\n",
    "Training with `WebDataset` can be carried out completely without local storage; this is the usual setup in the cloud\n",
    "and on high speed compute clusters. When running locally on a desktop, you may want to cache the data, and for that,\n",
    "you set a `cache_dir` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ced2827a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:20.499880Z",
     "iopub.status.busy": "2023-12-18T08:08:20.499604Z",
     "iopub.status.idle": "2023-12-18T08:08:22.678443Z",
     "shell.execute_reply": "2023-12-18T08:08:22.677488Z"
    },
    "papermill": {
     "duration": 2.184855,
     "end_time": "2023-12-18T08:08:22.681177",
     "exception": false,
     "start_time": "2023-12-18T08:08:20.496322",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.distributed as dist\n",
    "from torch.nn.parallel import DistributedDataParallel\n",
    "from torchvision.models import resnet50\n",
    "from torchvision import datasets, transforms\n",
    "import ray\n",
    "import webdataset as wds\n",
    "import dataclasses\n",
    "import time\n",
    "from collections import deque\n",
    "from typing import Optional\n",
    "\n",
    "\n",
    "def enumerate_report(seq, delta, growth=1.0):\n",
    "    last = 0\n",
    "    count = 0\n",
    "    for count, item in enumerate(seq):\n",
    "        now = time.time()\n",
    "        if now - last > delta:\n",
    "            last = now\n",
    "            yield count, item, True\n",
    "        else:\n",
    "            yield count, item, False\n",
    "        delta *= growth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "14f03e18",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:22.691338Z",
     "iopub.status.busy": "2023-12-18T08:08:22.691142Z",
     "iopub.status.idle": "2023-12-18T08:08:22.694517Z",
     "shell.execute_reply": "2023-12-18T08:08:22.693963Z"
    },
    "papermill": {
     "duration": 0.010232,
     "end_time": "2023-12-18T08:08:22.696429",
     "exception": false,
     "start_time": "2023-12-18T08:08:22.686197",
     "status": "completed"
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "epochs = 10\n",
    "maxsteps = int(1e12)\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "28aff135",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:22.704868Z",
     "iopub.status.busy": "2023-12-18T08:08:22.704681Z",
     "iopub.status.idle": "2023-12-18T08:08:22.707649Z",
     "shell.execute_reply": "2023-12-18T08:08:22.707100Z"
    },
    "papermill": {
     "duration": 0.00937,
     "end_time": "2023-12-18T08:08:22.709629",
     "exception": false,
     "start_time": "2023-12-18T08:08:22.700259",
     "status": "completed"
    },
    "tags": [
     "injected-parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "max_steps = 1000\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d23d386c",
   "metadata": {
    "papermill": {
     "duration": 0.003701,
     "end_time": "2023-12-18T08:08:22.717141",
     "exception": false,
     "start_time": "2023-12-18T08:08:22.713440",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Data Loading for Distributed Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d4278e4b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:22.725287Z",
     "iopub.status.busy": "2023-12-18T08:08:22.725082Z",
     "iopub.status.idle": "2023-12-18T08:08:22.728718Z",
     "shell.execute_reply": "2023-12-18T08:08:22.728072Z"
    },
    "papermill": {
     "duration": 0.009809,
     "end_time": "2023-12-18T08:08:22.730635",
     "exception": false,
     "start_time": "2023-12-18T08:08:22.720826",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Datasets are just collections of shards in the cloud. We usually specify\n",
    "# them using {lo..hi} brace notation (there is also a YAML spec for more complex\n",
    "# datasets).\n",
    "\n",
    "bucket = \"https://storage.googleapis.com/webdataset/fake-imagenet\"\n",
    "trainset_url = bucket + \"/imagenet-train-{000000..001281}.tar\"\n",
    "valset_url = bucket + \"/imagenet-val-{000000..000049}.tar\"\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7c6acd62",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:22.738384Z",
     "iopub.status.busy": "2023-12-18T08:08:22.738121Z",
     "iopub.status.idle": "2023-12-18T08:08:22.743963Z",
     "shell.execute_reply": "2023-12-18T08:08:22.743050Z"
    },
    "papermill": {
     "duration": 0.011895,
     "end_time": "2023-12-18T08:08:22.745927",
     "exception": false,
     "start_time": "2023-12-18T08:08:22.734032",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "not running in colab, caching data locally in ./_cache\n"
     ]
    }
   ],
   "source": [
    "# If running in the cloud or with a fast network storage system, we don't\n",
    "# need any local storage.\n",
    "\n",
    "if \"google.colab\" in sys.modules:\n",
    "    cache_dir = None\n",
    "    print(\"running on colab, streaming data directly from storage\")\n",
    "else:\n",
    "    cache_dir = \"./_cache\"\n",
    "    print(f\"not running in colab, caching data locally in {cache_dir}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "adfd84eb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:22.754286Z",
     "iopub.status.busy": "2023-12-18T08:08:22.753942Z",
     "iopub.status.idle": "2023-12-18T08:08:22.762838Z",
     "shell.execute_reply": "2023-12-18T08:08:22.761805Z"
    },
    "papermill": {
     "duration": 0.015494,
     "end_time": "2023-12-18T08:08:22.764925",
     "exception": false,
     "start_time": "2023-12-18T08:08:22.749431",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# The dataloader pipeline is a fairly typical `IterableDataset` pipeline\n",
    "# for PyTorch\n",
    "\n",
    "\n",
    "def make_dataloader_train():\n",
    "    \"\"\"Create a DataLoader for training on the ImageNet dataset using WebDataset.\"\"\"\n",
    "\n",
    "    transform = transforms.Compose(\n",
    "        [\n",
    "            transforms.RandomResizedCrop(224),\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            transforms.ToTensor(),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    def make_sample(sample):\n",
    "        return transform(sample[\"jpg\"]), sample[\"cls\"]\n",
    "\n",
    "    # This is the basic WebDataset definition: it starts with a URL and add shuffling,\n",
    "    # decoding, and augmentation. Note `resampled=True`; this is essential for\n",
    "    # distributed training to work correctly.\n",
    "    trainset = wds.WebDataset(trainset_url, resampled=True, cache_dir=cache_dir, nodesplitter=wds.split_by_node)\n",
    "    trainset = trainset.shuffle(1000).decode(\"pil\").map(make_sample)\n",
    "\n",
    "    # For IterableDataset objects, the batching needs to happen in the dataset.\n",
    "    trainset = trainset.batched(64)\n",
    "    trainloader = wds.WebLoader(trainset, batch_size=None, num_workers=4)\n",
    "\n",
    "    # We unbatch, shuffle, and rebatch to mix samples from different workers.\n",
    "    trainloader = trainloader.unbatched().shuffle(1000).batched(batch_size)\n",
    "\n",
    "    # A resampled dataset is infinite size, but we can recreate a fixed epoch length.\n",
    "    trainloader = trainloader.with_epoch(1282 * 100 // 64)\n",
    "\n",
    "    return trainloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "199d1b7a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:22.773486Z",
     "iopub.status.busy": "2023-12-18T08:08:22.773100Z",
     "iopub.status.idle": "2023-12-18T08:08:23.975916Z",
     "shell.execute_reply": "2023-12-18T08:08:23.975205Z"
    },
    "papermill": {
     "duration": 1.209917,
     "end_time": "2023-12-18T08:08:23.978341",
     "exception": false,
     "start_time": "2023-12-18T08:08:22.768424",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 3, 224, 224]) torch.Size([32])\n"
     ]
    }
   ],
   "source": [
    "# Let's try it out\n",
    "\n",
    "\n",
    "def make_dataloader(split=\"train\"):\n",
    "    \"\"\"Make a dataloader for training or validation.\"\"\"\n",
    "    if split == \"train\":\n",
    "        return make_dataloader_train()\n",
    "    elif split == \"val\":\n",
    "        return make_dataloader_val()  # not implemented for this notebook\n",
    "    else:\n",
    "        raise ValueError(f\"unknown split {split}\")\n",
    "\n",
    "\n",
    "# Try it out.\n",
    "os.environ[\"GOPEN_VERBOSE\"] = \"1\"\n",
    "sample = next(iter(make_dataloader()))\n",
    "print(sample[0].shape, sample[1].shape)\n",
    "os.environ[\"GOPEN_VERBOSE\"] = \"0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40ce90a0",
   "metadata": {
    "papermill": {
     "duration": 0.004566,
     "end_time": "2023-12-18T08:08:23.988045",
     "exception": false,
     "start_time": "2023-12-18T08:08:23.983479",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Standard PyTorch Training\n",
    "\n",
    "This is completely standard PyTorch training; nothing changes by using WebDataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e3bcfaee",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:23.999982Z",
     "iopub.status.busy": "2023-12-18T08:08:23.999533Z",
     "iopub.status.idle": "2023-12-18T08:08:24.009859Z",
     "shell.execute_reply": "2023-12-18T08:08:24.008598Z"
    },
    "papermill": {
     "duration": 0.019412,
     "end_time": "2023-12-18T08:08:24.012502",
     "exception": false,
     "start_time": "2023-12-18T08:08:23.993090",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# We gather all the configuration info into a single typed dataclass.\n",
    "\n",
    "\n",
    "@dataclasses.dataclass\n",
    "class Config:\n",
    "    epochs: int = 1\n",
    "    max_steps: int = int(1e18)\n",
    "    lr: float = 0.001\n",
    "    momentum: float = 0.9\n",
    "    rank: Optional[int] = None\n",
    "    world_size: int = 2\n",
    "    backend: str = \"nccl\"\n",
    "    master_addr: str = \"localhost\"\n",
    "    master_port: str = \"12355\"\n",
    "    report_s: float = 15.0\n",
    "    report_growth: float = 1.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3910e6a9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:24.024027Z",
     "iopub.status.busy": "2023-12-18T08:08:24.023530Z",
     "iopub.status.idle": "2023-12-18T08:08:24.040119Z",
     "shell.execute_reply": "2023-12-18T08:08:24.038717Z"
    },
    "papermill": {
     "duration": 0.026233,
     "end_time": "2023-12-18T08:08:24.043485",
     "exception": false,
     "start_time": "2023-12-18T08:08:24.017252",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train(config):\n",
    "    # Define the model, loss function, and optimizer\n",
    "    model = resnet50(pretrained=False).cuda()\n",
    "    if config.rank is not None:\n",
    "        model = DistributedDataParallel(model)\n",
    "    loss_fn = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=config.lr)\n",
    "\n",
    "    # Data loading code\n",
    "    trainloader = make_dataloader(split=\"train\")\n",
    "\n",
    "    losses, accuracies, steps = deque(maxlen=100), deque(maxlen=100), 0\n",
    "\n",
    "    # Training loop\n",
    "    for epoch in range(config.epochs):\n",
    "        for i, data, verbose in enumerate_report(trainloader, config.report_s):\n",
    "            inputs, labels = data[0].cuda(), data[1].cuda()\n",
    "\n",
    "            # zero the parameter gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # forward + backward + optimize\n",
    "            outputs = model(inputs)\n",
    "\n",
    "            # update statistics\n",
    "            loss = loss_fn(outputs, labels)\n",
    "            accuracy = (\n",
    "                (outputs.argmax(1) == labels).float().mean()\n",
    "            )  # calculate accuracy\n",
    "            losses.append(loss.item())\n",
    "            accuracies.append(accuracy.item())\n",
    "\n",
    "            if verbose and len(losses) > 0:\n",
    "                avgloss = sum(losses) / len(losses)\n",
    "                avgaccuracy = sum(accuracies) / len(accuracies)\n",
    "                print(\n",
    "                    f\"rank {config.rank} epoch {epoch:5d}/{i:9d} loss {avgloss:8.3f} acc {avgaccuracy:8.3f} {steps:9d}\",\n",
    "                    file=sys.stderr,\n",
    "                )\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            steps += len(labels)\n",
    "            if steps > config.max_steps:\n",
    "                print(\n",
    "                    \"finished training (max_steps)\",\n",
    "                    steps,\n",
    "                    config.max_steps,\n",
    "                    file=sys.stderr,\n",
    "                )\n",
    "                return\n",
    "\n",
    "    print(\"finished Training\", steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "66fe87f6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:24.058877Z",
     "iopub.status.busy": "2023-12-18T08:08:24.058369Z",
     "iopub.status.idle": "2023-12-18T08:08:31.854615Z",
     "shell.execute_reply": "2023-12-18T08:08:31.853761Z"
    },
    "papermill": {
     "duration": 7.807686,
     "end_time": "2023-12-18T08:08:31.857838",
     "exception": false,
     "start_time": "2023-12-18T08:08:24.050152",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tmb/proj/webdataset/venv/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/home/tmb/proj/webdataset/venv/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=None`.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "rank None epoch     0/        0 loss    7.259 acc    0.000         0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "finished training (max_steps) 1024 1000\n"
     ]
    }
   ],
   "source": [
    "# A quick smoke test of the training function.\n",
    "\n",
    "config = Config()\n",
    "config.epochs = 1\n",
    "config.max_steps = 1000\n",
    "train(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a5d346e",
   "metadata": {
    "papermill": {
     "duration": 0.007157,
     "end_time": "2023-12-18T08:08:31.874228",
     "exception": false,
     "start_time": "2023-12-18T08:08:31.867071",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Setting up Distributed Training with Ray\n",
    "\n",
    "Ray is a convenient distributed computing framework. We are using it here to start up the training\n",
    "jobs on multiple GPUs. You can use `torch.distributed.launch` or other such tools as well with the above\n",
    "code. Ray has the advantage that it is runtime environment independent; you set up your Ray cluster\n",
    "in whatever way works for your environment, and afterwards, this code will run in it without change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "62b3afab",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:31.890160Z",
     "iopub.status.busy": "2023-12-18T08:08:31.889824Z",
     "iopub.status.idle": "2023-12-18T08:08:31.894712Z",
     "shell.execute_reply": "2023-12-18T08:08:31.894135Z"
    },
    "papermill": {
     "duration": 0.016636,
     "end_time": "2023-12-18T08:08:31.898034",
     "exception": false,
     "start_time": "2023-12-18T08:08:31.881398",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "@ray.remote(num_gpus=1)\n",
    "def train_on_ray(rank, config):\n",
    "    \"\"\"Set up distributed torch env and train the model on this node.\"\"\"\n",
    "    # Set up distributed PyTorch.\n",
    "    if rank is not None:\n",
    "        os.environ[\"MASTER_ADDR\"] = config.master_addr\n",
    "        os.environ[\"MASTER_PORT\"] = config.master_port\n",
    "        dist.init_process_group(\n",
    "            backend=config.backend, rank=rank, world_size=config.world_size\n",
    "        )\n",
    "        config.rank = rank\n",
    "        # Ray will automatically set CUDA_VISIBLE_DEVICES for each task.\n",
    "    train(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c0886142",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:08:31.914436Z",
     "iopub.status.busy": "2023-12-18T08:08:31.914208Z",
     "iopub.status.idle": "2023-12-18T08:09:03.108011Z",
     "shell.execute_reply": "2023-12-18T08:09:03.107406Z"
    },
    "papermill": {
     "duration": 31.203907,
     "end_time": "2023-12-18T08:09:03.109358",
     "exception": false,
     "start_time": "2023-12-18T08:08:31.905451",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-12-18 00:08:34,073\tINFO worker.py:1664 -- Started a local Ray instance. View the dashboard at \u001b[1m\u001b[32m127.0.0.1:8265 \u001b[39m\u001b[22m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Config(epochs=10, max_steps=1000, lr=0.001, momentum=0.9, rank=None, world_size=2, backend='nccl', master_addr='localhost', master_port='12355', report_s=15.0, report_growth=1.1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(train_on_ray pid=715298)\u001b[0m /home/tmb/proj/webdataset/venv/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "\u001b[36m(train_on_ray pid=715298)\u001b[0m   warnings.warn(\n",
      "\u001b[36m(train_on_ray pid=715298)\u001b[0m /home/tmb/proj/webdataset/venv/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=None`.\n",
      "\u001b[36m(train_on_ray pid=715298)\u001b[0m   warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(train_on_ray pid=715299)\u001b[0m rank 0 epoch     0/        0 loss    6.938 acc    0.000         0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(train_on_ray pid=715299)\u001b[0m /home/tmb/proj/webdataset/venv/lib/python3.10/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "\u001b[36m(train_on_ray pid=715299)\u001b[0m   warnings.warn(\n",
      "\u001b[36m(train_on_ray pid=715299)\u001b[0m /home/tmb/proj/webdataset/venv/lib/python3.10/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=None`.\n",
      "\u001b[36m(train_on_ray pid=715299)\u001b[0m   warnings.warn(msg)\n",
      "\u001b[36m(train_on_ray pid=715298)\u001b[0m rank 1 epoch     0/       23 loss    5.942 acc    0.053       736\u001b[32m [repeated 2x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/ray-logging.html#log-deduplication for more options.)\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(train_on_ray pid=715298)\u001b[0m finished training (max_steps) 1024 1000\n",
      "\u001b[36m(train_on_ray pid=715299)\u001b[0m rank 0 epoch     0/       23 loss    6.005 acc    0.033       736\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[None, None]\n"
     ]
    }
   ],
   "source": [
    "if not ray.is_initialized():\n",
    "    ray.init()\n",
    "\n",
    "ray.available_resources()[\"GPU\"]\n",
    "\n",
    "\n",
    "def distributed_training(config):\n",
    "    \"\"\"Perform distributed training with the given config.\"\"\"\n",
    "    num_gpus = ray.available_resources()[\"GPU\"]\n",
    "    config.world_size = min(config.world_size, num_gpus)\n",
    "    results = ray.get(\n",
    "        [train_on_ray.remote(i, config) for i in range(config.world_size)]\n",
    "    )\n",
    "    print(results)\n",
    "\n",
    "\n",
    "config = Config()\n",
    "config.epochs = epochs\n",
    "config.max_steps = max_steps\n",
    "config.batch_size = batch_size\n",
    "print(config)\n",
    "distributed_training(config)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.10.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 48.73453,
   "end_time": "2023-12-18T08:09:08.334027",
   "environment_variables": {},
   "exception": null,
   "input_path": "train-resnet50-multiray-wds.ipynb",
   "output_path": "out/_train-resnet50-multiray-wds.ipynb",
   "parameters": {
    "max_steps": 1000
   },
   "start_time": "2023-12-18T08:08:19.599497",
   "version": "2.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}