{ "cells": [ { "cell_type": "markdown", "id": "independent-thread", "metadata": {}, "source": [ "---\n", "title: \"Maggy distributed training ResNet-50 on ImageNet (Petastorm)\"\n", "date: 2021-05-03\n", "type: technical_note\n", "draft: false\n", "---" ] }, { "cell_type": "markdown", "id": "therapeutic-working", "metadata": {}, "source": [ "## Training ResNet-50 on ImageNet from a Petastorm dataset\n", "In this notebook, we are going to train a ResNet-50 network on a subset of 10 labels of the original ImageNet dataset. In order to improve our I/O time compared to the standard ImageNet training, we are going to use the Petastorm version of the dataset created in `ImageNet_to_petastorm`." ] }, { "cell_type": "code", "execution_count": 1, "id": "parallel-craft", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Starting Spark application\n" ] }, { "data": { "text/html": [ "\n", "
IDYARN Application IDKindStateSpark UIDriver log
186application_1617699042861_0013pysparkidleLinkLink
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "SparkSession available as 'spark'.\n" ] } ], "source": [ "from hops import hdfs\n", "from torchvision import models" ] }, { "cell_type": "markdown", "id": "recent-quest", "metadata": {}, "source": [ "### Defining the training function\n", "Just as in the notebooks before, we define our training function. Instead of using the default PyTorch dataloader however, we use the MaggyPetastormDataLoader. This DataLoader aims to mimic its PyTorch counterpart as close as possible. One key difference is that transforms have to be passed to the DataLoader instead of the dataset. You can do so with the `transform_spec`." ] }, { "cell_type": "code", "execution_count": 2, "id": "buried-chamber", "metadata": {}, "outputs": [], "source": [ "def train_fn(module, hparams, train_set, test_set):\n", " \n", " import time \n", " import torch\n", " from torchvision import transforms as T\n", " \n", " from maggy.core.patching import MaggyPetastormDataLoader\n", "\n", " \n", " model = module(**hparams)\n", " \n", " n_epochs = 10\n", " batch_size = 64\n", " lr_base = 0.1 * 2*batch_size/256\n", " \n", " def transform(image_net_row):\n", " transform = T.Compose([\n", " T.ToTensor(),\n", " T.RandomHorizontalFlip()\n", " ])\n", " return {\"image\": transform(image_net_row['image']), \"label\": image_net_row['label']}\n", "\n", " # Parameters as in https://arxiv.org/pdf/1706.02677.pdf\n", " optimizer = torch.optim.SGD(model.parameters(), lr=lr_base, momentum=0.9, weight_decay=0.0001, nesterov=True)\n", " loss_criterion = torch.nn.CrossEntropyLoss()\n", " \n", " train_loader = MaggyPetastormDataLoader(train_set, batch_size=batch_size, transform_spec=transform)\n", " test_loader = MaggyPetastormDataLoader(test_set, batch_size=batch_size, transform_spec=transform)\n", " \n", " def eval_model(model, test_loader):\n", " acc = 0\n", " model.eval()\n", " img_cnt = 0\n", " with torch.no_grad():\n", " for data in test_loader:\n", " img, label = data[\"image\"].float(), data[\"label\"].float()\n", " prediction = model(img)\n", " acc += torch.sum(torch.argmax(prediction, dim=1) == label).detach()\n", " img_cnt += len(label.detach())\n", " acc = acc/float(img_cnt)\n", " print(\"Test accuracy: {:.3f}\\n\".format(acc) + 20*\"-\")\n", " return acc\n", "\n", " model.train()\n", " t_0 = time.time()\n", " for epoch in range(n_epochs):\n", " print(\"-\"*20 + \"\\nStarting new epoch\\n\")\n", " model.train()\n", " t_start = time.time()\n", " for idx, data in enumerate(train_loader):\n", " optimizer.zero_grad()\n", " img, label = data[\"image\"].float(), data[\"label\"].float()\n", " prediction = model(img)\n", " loss = loss_criterion(prediction, label.long())\n", " loss.backward()\n", " optimizer.step()\n", " if idx%10 == 0:\n", " print(f\"Working on batch {idx}.\")\n", " t_end = time.time()\n", " print(\"Epoch training took {:.0f}s.\\n\".format(t_end - t_start))\n", " acc = eval_model(model, test_loader)\n", " t_1 = time.time()\n", " minutes, seconds = divmod(t_1 - t_0, 60)\n", " hours, minutes = divmod(minutes, 60)\n", " print(\"-\"*20 + \"\\nTotal training time: {:.0f}h {:.0f}m {:.0f}s.\".format(hours, minutes, seconds))\n", " return float(acc)" ] }, { "cell_type": "code", "execution_count": 3, "id": "advised-gospel", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "True True" ] } ], "source": [ "train_ds = hdfs.project_path() + \"DataSets/ImageNet/PetastormImageNette/train\"\n", "test_ds = hdfs.project_path() + \"DataSets/ImageNet/PetastormImageNette/test\"\n", "print(hdfs.exists(train_ds), hdfs.exists(test_ds))" ] }, { "cell_type": "markdown", "id": "aerial-acting", "metadata": {}, "source": [ "### Configuring the experiment\n", "In this example we are using the PyTorch provided implementation of ResNet50. We therefore do not need to define our own module. In the hparams argument, we can pass any arguments for the network. Of course, this is more a convenience mechanism than a necessity. You could also just define them in the training function itself. However, passing them in the config has the advantage that you can automate distributed training after e.g. hyperparameter search." ] }, { "cell_type": "code", "execution_count": 4, "id": "consistent-cambodia", "metadata": {}, "outputs": [], "source": [ "from maggy import experiment\n", "from maggy.experiment_config import TorchDistributedConfig\n", "\n", "config = TorchDistributedConfig(name='ImageNet_petastorm', module=models.resnet50, hparams={\"pretrained\": False}, train_set=train_ds, test_set=test_ds, backend=\"torch\")" ] }, { "cell_type": "code", "execution_count": null, "id": "nasty-wound", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "b99b8c27640544a4981d749be29d262e", "version_major": 2, "version_minor": 0 }, "text/plain": [ "HBox(children=(FloatProgress(value=0.0, description='Maggy experiment', max=1.0, style=ProgressStyle(descripti…" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "0: Awaiting worker reservations.\n", "1: Awaiting worker reservations.\n", "0: All executors registered: True\n", "0: Reservations complete, configuring PyTorch.\n", "0: Torch config is {'MASTER_ADDR': '10.0.0.4', 'MASTER_PORT': '46351', 'WORLD_SIZE': '2', 'RANK': '0', 'LOCAL_RANK': '0', 'NCCL_BLOCKING_WAIT': '1', 'NCCL_DEBUG': 'INFO'}\n", "0: Starting distributed training.\n", "1: All executors registered: True\n", "1: Reservations complete, configuring PyTorch.\n", "1: Torch config is {'MASTER_ADDR': '10.0.0.4', 'MASTER_PORT': '46351', 'WORLD_SIZE': '2', 'RANK': '1', 'LOCAL_RANK': '0', 'NCCL_BLOCKING_WAIT': '1', 'NCCL_DEBUG': 'INFO'}\n", "1: Starting distributed training.\n", "0: Petastorm dataset detected in folder hdfs://rpc.namenode.service.consul:8020/Projects/PyTorch_spark_minimal/DataSets/ImageNet/PetastormImageNette/train\n", "0: Petastorm dataset detected in folder hdfs://rpc.namenode.service.consul:8020/Projects/PyTorch_spark_minimal/DataSets/ImageNet/PetastormImageNette/test\n", "1: Petastorm dataset detected in folder hdfs://rpc.namenode.service.consul:8020/Projects/PyTorch_spark_minimal/DataSets/ImageNet/PetastormImageNette/train\n", "1: Petastorm dataset detected in folder hdfs://rpc.namenode.service.consul:8020/Projects/PyTorch_spark_minimal/DataSets/ImageNet/PetastormImageNette/test\n", "0: --------------------\n", "Starting new epoch\n", "\n", "1: --------------------\n", "Starting new epoch\n", "\n", "0: Working on batch 0.\n", "1: Working on batch 0.\n", "0: Working on batch 10.\n", "1: Working on batch 10.\n", "0: Working on batch 20.\n", "1: Working on batch 20.\n", "0: Working on batch 30.\n", "1: Working on batch 30.\n", "0: Working on batch 40.\n", "1: Working on batch 40.\n", "0: Working on batch 50.\n", "1: Working on batch 50.\n", "0: Working on batch 60.\n", "1: Working on batch 60.\n", "0: Working on batch 70.\n", "1: Working on batch 70.\n", "0: Epoch training took 149s.\n", "\n", "1: Epoch training took 149s.\n", "\n", "0: Test accuracy: 0.134\n", "--------------------\n", "0: --------------------\n", "Starting new epoch\n", "\n", "1: Test accuracy: 0.093\n", "--------------------\n", "1: --------------------\n", "Starting new epoch\n", "\n", "1: Working on batch 0.\n", "0: Working on batch 0.\n", "1: Working on batch 10.\n", "0: Working on batch 10.\n", "1: Working on batch 20.\n", "0: Working on batch 20.\n", "1: Working on batch 30.\n", "0: Working on batch 30.\n", "1: Working on batch 40.\n", "0: Working on batch 40.\n", "0: Working on batch 50.\n", "1: Working on batch 50.\n", "1: Working on batch 60.\n", "0: Working on batch 60.\n", "1: Working on batch 70.\n", "0: Working on batch 70.\n", "1: Epoch training took 149s.\n", "\n", "0: Epoch training took 150s.\n", "\n", "0: Test accuracy: 0.171\n", "--------------------\n", "0: --------------------\n", "Starting new epoch\n", "\n", "1: Test accuracy: 0.126\n", "--------------------\n", "1: --------------------\n", "Starting new epoch\n", "\n", "0: Working on batch 0.\n", "1: Working on batch 0.\n", "0: Working on batch 10.\n", "1: Working on batch 10.\n", "1: Working on batch 20.\n", "0: Working on batch 20.\n", "0: Working on batch 30.\n", "1: Working on batch 30.\n", "1: Working on batch 40.\n", "0: Working on batch 40.\n", "0: Working on batch 50.\n", "1: Working on batch 50.\n", "0: Working on batch 60.\n", "1: Working on batch 60.\n", "1: Working on batch 70.\n", "0: Working on batch 70.\n", "1: Epoch training took 150s.\n", "\n", "0: Epoch training took 150s.\n", "\n", "0: Test accuracy: 0.235\n", "--------------------\n", "0: --------------------\n", "Starting new epoch\n", "\n", "1: Test accuracy: 0.192\n", "--------------------\n", "1: --------------------\n", "Starting new epoch\n", "\n", "1: Working on batch 0.\n", "0: Working on batch 0.\n", "0: Working on batch 10.\n", "1: Working on batch 10.\n", "1: Working on batch 20.\n", "0: Working on batch 20.\n", "0: Working on batch 30.\n", "1: Working on batch 30.\n", "0: Working on batch 40.\n", "1: Working on batch 40.\n", "0: Working on batch 50.\n", "1: Working on batch 50.\n", "0: Working on batch 60.\n", "1: Working on batch 60.\n", "1: Working on batch 70.\n", "0: Working on batch 70.\n", "0: Epoch training took 150s.\n", "\n", "1: Epoch training took 149s.\n", "\n", "1: Test accuracy: 0.277\n", "--------------------\n", "1: --------------------\n", "Starting new epoch\n", "\n", "0: Test accuracy: 0.188\n", "--------------------\n", "0: --------------------\n", "Starting new epoch\n", "\n", "1: Working on batch 0.\n", "0: Working on batch 0.\n", "1: Working on batch 10.\n", "0: Working on batch 10.\n", "1: Working on batch 20.\n", "0: Working on batch 20.\n", "1: Working on batch 30.\n", "0: Working on batch 30.\n", "1: Working on batch 40.\n", "0: Working on batch 40.\n", "1: Working on batch 50.\n", "0: Working on batch 50.\n", "1: Working on batch 60.\n", "0: Working on batch 60.\n", "1: Working on batch 70.\n", "0: Working on batch 70.\n", "1: Epoch training took 150s.\n", "\n", "0: Epoch training took 151s.\n", "\n", "0: Test accuracy: 0.305\n", "--------------------\n", "0: --------------------\n", "Starting new epoch\n", "\n", "1: Test accuracy: 0.311\n", "--------------------\n", "1: --------------------\n", "Starting new epoch\n", "\n", "0: Working on batch 0.\n", "1: Working on batch 0.\n", "1: Working on batch 10.\n", "0: Working on batch 10.\n", "0: Working on batch 20.\n", "1: Working on batch 20.\n", "1: Working on batch 30.\n", "0: Working on batch 30.\n", "1: Working on batch 40.\n", "0: Working on batch 40.\n", "0: Working on batch 50.\n", "1: Working on batch 50.\n", "1: Working on batch 60.\n", "0: Working on batch 60.\n", "0: Working on batch 70.\n", "1: Working on batch 70.\n", "0: Epoch training took 151s.\n", "\n", "1: Epoch training took 150s.\n", "\n", "0: Test accuracy: 0.377\n", "--------------------\n", "0: --------------------\n", "Starting new epoch\n", "\n", "1: Test accuracy: 0.335\n", "--------------------\n", "1: --------------------\n", "Starting new epoch\n", "\n", "0: Working on batch 0.\n", "1: Working on batch 0.\n", "0: Working on batch 10.\n", "1: Working on batch 10.\n", "0: Working on batch 20.\n", "1: Working on batch 20.\n", "0: Working on batch 30.\n", "1: Working on batch 30.\n", "0: Working on batch 40.\n", "1: Working on batch 40.\n" ] } ], "source": [ "result = experiment.lagom(train_fn, config)" ] }, { "cell_type": "code", "execution_count": null, "id": "found-proposition", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "PySpark", "language": "python", "name": "pysparkkernel" }, "language_info": { "codemirror_mode": { "name": "python", "version": 3 }, "mimetype": "text/x-python", "name": "pyspark", "pygments_lexer": "python3" } }, "nbformat": 4, "nbformat_minor": 5 }