examples/out/train-resnet50-wds.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "08d891b9",
   "metadata": {
    "papermill": {
     "duration": 0.00308,
     "end_time": "2023-12-18T08:05:54.086460",
     "exception": false,
     "start_time": "2023-12-18T08:05:54.083380",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Resnet 50 Training on (Fake)Imagenet with WebDataset\n",
    "\n",
    "This notebook illustrates how to use WebDataset with PyTorch training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e26d2ea8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:05:54.093086Z",
     "iopub.status.busy": "2023-12-18T08:05:54.092820Z",
     "iopub.status.idle": "2023-12-18T08:05:56.290837Z",
     "shell.execute_reply": "2023-12-18T08:05:56.290055Z"
    },
    "papermill": {
     "duration": 2.204807,
     "end_time": "2023-12-18T08:05:56.294017",
     "exception": false,
     "start_time": "2023-12-18T08:05:54.089210",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "from functools import partial\n",
    "from pprint import pprint\n",
    "import random\n",
    "from collections import deque\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import sys\n",
    "import os\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.models import resnet50\n",
    "from torch.utils.data import DataLoader\n",
    "from torch import nn, optim\n",
    "\n",
    "# helpers\n",
    "\n",
    "import time\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": "12d63615",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:05:56.306788Z",
     "iopub.status.busy": "2023-12-18T08:05:56.306504Z",
     "iopub.status.idle": "2023-12-18T08:05:56.315465Z",
     "shell.execute_reply": "2023-12-18T08:05:56.314880Z"
    },
    "papermill": {
     "duration": 0.018343,
     "end_time": "2023-12-18T08:05:56.318437",
     "exception": false,
     "start_time": "2023-12-18T08:05:56.300094",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# We usually abbreviate webdataset as wds\n",
    "import webdataset as wds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8cfc1a4c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:05:56.330751Z",
     "iopub.status.busy": "2023-12-18T08:05:56.330547Z",
     "iopub.status.idle": "2023-12-18T08:05:56.334271Z",
     "shell.execute_reply": "2023-12-18T08:05:56.333716Z"
    },
    "papermill": {
     "duration": 0.013444,
     "end_time": "2023-12-18T08:05:56.337364",
     "exception": false,
     "start_time": "2023-12-18T08:05:56.323920",
     "status": "completed"
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# parameters\n",
    "epochs = 1\n",
    "max_steps = int(1e12)\n",
    "batchsize = 32\n",
    "bucket = \"https://storage.googleapis.com/webdataset/fake-imagenet\"\n",
    "training_urls = bucket + \"/imagenet-train-{000000..001281}.tar\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ce1632db",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:05:56.349755Z",
     "iopub.status.busy": "2023-12-18T08:05:56.349538Z",
     "iopub.status.idle": "2023-12-18T08:05:56.353293Z",
     "shell.execute_reply": "2023-12-18T08:05:56.352259Z"
    },
    "papermill": {
     "duration": 0.013227,
     "end_time": "2023-12-18T08:05:56.356229",
     "exception": false,
     "start_time": "2023-12-18T08:05:56.343002",
     "status": "completed"
    },
    "tags": [
     "injected-parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "max_steps = 10000\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "246e5ad0",
   "metadata": {
    "papermill": {
     "duration": 0.0055,
     "end_time": "2023-12-18T08:05:56.367334",
     "exception": false,
     "start_time": "2023-12-18T08:05:56.361834",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Data Loader Construction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "251135b9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:05:56.379612Z",
     "iopub.status.busy": "2023-12-18T08:05:56.379190Z",
     "iopub.status.idle": "2023-12-18T08:05:56.505443Z",
     "shell.execute_reply": "2023-12-18T08:05:56.503673Z"
    },
    "papermill": {
     "duration": 0.135855,
     "end_time": "2023-12-18T08:05:56.508603",
     "exception": false,
     "start_time": "2023-12-18T08:05:56.372748",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "not running in colab, caching data locally in ./_cache\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# WebDataset is designed to work without any local storage. Use caching\n",
    "# only if you are on a desktop with slow networking.\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",
    "    !mkdir -p ./_cache\n",
    "    cache_dir = \"./_cache\"\n",
    "    print(f\"not running in colab, caching data locally in {cache_dir}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "78ee073f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:05:56.521907Z",
     "iopub.status.busy": "2023-12-18T08:05:56.521344Z",
     "iopub.status.idle": "2023-12-18T08:05:56.536947Z",
     "shell.execute_reply": "2023-12-18T08:05:56.535632Z"
    },
    "papermill": {
     "duration": 0.025966,
     "end_time": "2023-12-18T08:05:56.540257",
     "exception": false,
     "start_time": "2023-12-18T08:05:56.514291",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# The standard TorchVision transformations.\n",
    "\n",
    "transform_train = transforms.Compose(\n",
    "    [\n",
    "        transforms.RandomResizedCrop(224),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "def make_sample(sample, val=False):\n",
    "    \"\"\"Take a decoded sample dictionary, augment it, and return an (image, label) tuple.\"\"\"\n",
    "    assert not val, \"only implemented training dataset for this notebook\"\n",
    "    image = sample[\"jpg\"]\n",
    "    label = sample[\"cls\"]\n",
    "    return transform_train(image), label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3cade01b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:05:56.552557Z",
     "iopub.status.busy": "2023-12-18T08:05:56.552035Z",
     "iopub.status.idle": "2023-12-18T08:05:56.571288Z",
     "shell.execute_reply": "2023-12-18T08:05:56.569800Z"
    },
    "papermill": {
     "duration": 0.029497,
     "end_time": "2023-12-18T08:05:56.574772",
     "exception": false,
     "start_time": "2023-12-18T08:05:56.545275",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create the datasets with shard and sample shuffling and decoding.\n",
    "trainset = wds.WebDataset(\n",
    "    training_urls, resampled=True, cache_dir=cache_dir, shardshuffle=True\n",
    ")\n",
    "trainset = trainset.shuffle(1000).decode(\"pil\").map(make_sample)\n",
    "\n",
    "# Since this is an IterableDataset, PyTorch requires that we batch in the dataset.\n",
    "# WebLoader is PyTorch DataLoader with some convenience methods.\n",
    "trainset = trainset.batched(64)\n",
    "trainloader = wds.WebLoader(trainset, batch_size=None, num_workers=4)\n",
    "\n",
    "# Unbatch, shuffle between workers, then rebatch.\n",
    "trainloader = trainloader.unbatched().shuffle(1000).batched(64)\n",
    "\n",
    "# Since we are using resampling, the dataset is infinite; set an artificial epoch size.\n",
    "trainloader = trainloader.with_epoch(1282 * 100 // 64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f704fae9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:05:56.589035Z",
     "iopub.status.busy": "2023-12-18T08:05:56.588498Z",
     "iopub.status.idle": "2023-12-18T08:05:57.784084Z",
     "shell.execute_reply": "2023-12-18T08:05:57.783281Z"
    },
    "papermill": {
     "duration": 1.205681,
     "end_time": "2023-12-18T08:05:57.786404",
     "exception": false,
     "start_time": "2023-12-18T08:05:56.580723",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 3, 224, 224]) torch.Size([64])\n"
     ]
    }
   ],
   "source": [
    "# Smoke test it.\n",
    "\n",
    "os.environ[\"GOPEN_VERBOSE\"] = \"1\"\n",
    "images, classes = next(iter(trainloader))\n",
    "print(images.shape, classes.shape)\n",
    "os.environ[\"GOPEN_VERBOSE\"] = \"0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d9b6fc7",
   "metadata": {
    "papermill": {
     "duration": 0.003978,
     "end_time": "2023-12-18T08:05:57.794993",
     "exception": false,
     "start_time": "2023-12-18T08:05:57.791015",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# PyTorch Training\n",
    "\n",
    "This is a typical PyTorch training pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "00e1dcbc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:05:57.805267Z",
     "iopub.status.busy": "2023-12-18T08:05:57.804797Z",
     "iopub.status.idle": "2023-12-18T08:05:59.863245Z",
     "shell.execute_reply": "2023-12-18T08:05:59.862319Z"
    },
    "papermill": {
     "duration": 2.067737,
     "end_time": "2023-12-18T08:05:59.866647",
     "exception": false,
     "start_time": "2023-12-18T08:05:57.798910",
     "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"
     ]
    }
   ],
   "source": [
    "# The usual PyTorch model definition. We use an uninitialized ResNet50 model.\n",
    "\n",
    "model = resnet50(pretrained=False)\n",
    "\n",
    "# Define the loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)\n",
    "\n",
    "# Move the model to the GPU if available\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dac1d651",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T08:05:59.881221Z",
     "iopub.status.busy": "2023-12-18T08:05:59.881022Z",
     "iopub.status.idle": "2023-12-18T08:07:03.479174Z",
     "shell.execute_reply": "2023-12-18T08:07:03.478131Z"
    },
    "papermill": {
     "duration": 63.608914,
     "end_time": "2023-12-18T08:07:03.482608",
     "exception": false,
     "start_time": "2023-12-18T08:05:59.873694",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,     8] loss: 5.08721 correct: 0.05664\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,    22] loss: 6.45160 correct: 0.07599\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,    36] loss: 5.62036 correct: 0.07986\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,    50] loss: 5.01627 correct: 0.08687\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,    64] loss: 4.59824 correct: 0.09204\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,    78] loss: 4.26513 correct: 0.10958\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,    92] loss: 4.03628 correct: 0.12211\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,   105] loss: 3.77721 correct: 0.14188\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,   118] loss: 3.19511 correct: 0.16109\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,   131] loss: 2.86325 correct: 0.18547\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,   144] loss: 2.68119 correct: 0.21000\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,   157] loss: 2.54697 correct: 0.24156\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished Training\n"
     ]
    }
   ],
   "source": [
    "losses, accuracies = deque(maxlen=100), deque(maxlen=100)\n",
    "\n",
    "steps = 0\n",
    "\n",
    "# Train the model\n",
    "for epoch in range(epochs):\n",
    "    for i, data, verbose in enumerate_report(trainloader, 5):\n",
    "        # get the inputs; data is a list of [inputs, labels]\n",
    "        inputs, labels = data[0].to(device), data[1].to(device)\n",
    "\n",
    "        # zero the parameter gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # forward + backward + optimize\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        pred = outputs.cpu().detach().argmax(dim=1, keepdim=True)\n",
    "        correct = pred.eq(labels.cpu().view_as(pred)).sum().item()\n",
    "        accuracy = correct / float(len(labels))\n",
    "\n",
    "        losses.append(loss.item())\n",
    "        accuracies.append(accuracy)\n",
    "        steps += len(inputs)\n",
    "\n",
    "        if verbose and len(losses) > 5:\n",
    "            print(\n",
    "                \"[%d, %5d] loss: %.5f correct: %.5f\"\n",
    "                % (epoch + 1, i + 1, np.mean(losses), np.mean(accuracies))\n",
    "            )\n",
    "            running_loss = 0.0\n",
    "\n",
    "        if steps > max_steps:\n",
    "            break\n",
    "\n",
    "    if steps > max_steps:\n",
    "        break\n",
    "\n",
    "print(\"Finished Training\")"
   ]
  }
 ],
 "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": 71.411865,
   "end_time": "2023-12-18T08:07:04.610424",
   "environment_variables": {},
   "exception": null,
   "input_path": "train-resnet50-wds.ipynb",
   "output_path": "out/_train-resnet50-wds.ipynb",
   "parameters": {
    "max_steps": 10000
   },
   "start_time": "2023-12-18T08:05:53.198559",
   "version": "2.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}