examples/notebooks/DistributedTripletMarginLossMNIST.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0SEqLaxOq7kw",
        "outputId": "55af6bc9-9c51-400d-dffc-931f384f57da"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting pytorch-metric-learning\n",
            "  Downloading pytorch_metric_learning-1.0.0-py3-none-any.whl (102 kB)\n",
            "\u001b[K     |████████████████████████████████| 102 kB 5.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from pytorch-metric-learning) (4.62.3)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.7/dist-packages (from pytorch-metric-learning) (1.0.1)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.7/dist-packages (from pytorch-metric-learning) (0.11.1+cu111)\n",
            "Requirement already satisfied: torch>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from pytorch-metric-learning) (1.10.0+cu111)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from pytorch-metric-learning) (1.19.5)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.6.0->pytorch-metric-learning) (3.10.0.2)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->pytorch-metric-learning) (3.0.0)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->pytorch-metric-learning) (1.1.0)\n",
            "Requirement already satisfied: scipy>=1.1.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->pytorch-metric-learning) (1.4.1)\n",
            "Requirement already satisfied: pillow!=8.3.0,>=5.3.0 in /usr/local/lib/python3.7/dist-packages (from torchvision->pytorch-metric-learning) (7.1.2)\n",
            "Installing collected packages: pytorch-metric-learning\n",
            "Successfully installed pytorch-metric-learning-1.0.0\n",
            "Collecting faiss-cpu\n",
            "  Downloading faiss_cpu-1.7.1.post2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (8.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 8.4 MB 5.2 MB/s \n",
            "\u001b[?25hInstalling collected packages: faiss-cpu\n",
            "Successfully installed faiss-cpu-1.7.1.post2\n"
          ]
        }
      ],
      "source": [
        "!pip install pytorch-metric-learning\n",
        "!pip install faiss-cpu # we're using cpu for this example"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000,
          "referenced_widgets": [
            "26e7efdf822b42d99290b2b8f2d7c050",
            "2a1a0715c79e473291feeecf52ede298",
            "48d0cf3885884015b45c6c471d0e2b1b",
            "aaa2a088e7de420f8bd6967033c87777",
            "d6e739eee399488aaf1e3c627e1f3221",
            "257593fba66f47ef8ecf23b3b5528b07",
            "9b861d5a4fa545a59ab089947c0f379e",
            "95e9f56ee5494094bb7780e426adc01d",
            "4a9be645d30545dc813dc71bc5eca1a3",
            "6341bdeb007b449b9876cce5de240c92",
            "39147f4e534040278af7f4be263da703",
            "0107a7e435ff4f06a912b98e190c2c9d",
            "82a231510951443a882e0a4d7e942467",
            "e626e3253ba1417596f6b5d3f87d8ebe",
            "fc0e2d892d8c4779aaefc3660f53670e",
            "6cb32fa30a4a4900a3d20b8ab3bdfe51",
            "ea267ba5a5694e85a68dbdd76f84ad5d",
            "c7f75365c3a54961a5698f0d8268e6a9",
            "48ea2d00bbaa46b7a52e4b3deed318ce",
            "0ddb2feda3c244439acda50c16d87732",
            "96cda59489624e21b67628b44a184e8b",
            "6fad8f3ff1884babbe4db769db6bdeae",
            "4c9ff1b060f34471a37d7ac5d19596fe",
            "ca98a6c54ef4473cb2b242d35ccc9fb6",
            "9fc0120ee9a94a09baedfeac38d0d9ff",
            "8b528385a13e4e3aa91bf5fcffce3409",
            "e372ec3e6cc44b5190241d47089d346c",
            "22c6745c4d474207b36c4c3acea542ce",
            "38653e48529a475eb0756bccb5820ce6",
            "6c4a3be721784447aa417eb861ed43af",
            "a59365e2120c4b098c97c27742dc8de6",
            "43de5b6f54dd4e69aaab95346425fa4d",
            "5d8e3ca6aec247528f8c67cd1efc7ffc",
            "7574947398c443a9b0c71a59e5a8598b",
            "25409ddba9e54cea97b43de0e4a45947",
            "a42ca39a12e44431a6ac7ae85a2e572d",
            "339792cd03f2408b9310baae564127fe",
            "91b6ab0c94194ca5a026170e703541b0",
            "787746306a4248f6a995e215d68f2e2c",
            "da95ae822e624f0897d8c89ea77e3277",
            "30c1a16ac224445fadcc2d1f7f141936",
            "d38ad67181854718b18b0a88228743aa",
            "ce1056721ac04f1bb851a177501278eb",
            "8dc29e4ab3bd4a6b8bf05270a206bb6e"
          ]
        },
        "id": "vYmfxuw1tYdm",
        "outputId": "3b5db4cd-13fc-4081-b0bf-5efc99d5ffbd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./data/MNIST/raw/train-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "26e7efdf822b42d99290b2b8f2d7c050",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/9912422 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Extracting ./data/MNIST/raw/train-images-idx3-ubyte.gz to ./data/MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./data/MNIST/raw/train-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "0107a7e435ff4f06a912b98e190c2c9d",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/28881 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Extracting ./data/MNIST/raw/train-labels-idx1-ubyte.gz to ./data/MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw/t10k-images-idx3-ubyte.gz\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "4c9ff1b060f34471a37d7ac5d19596fe",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/1648877 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Extracting ./data/MNIST/raw/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw\n",
            "\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
            "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "7574947398c443a9b0c71a59e5a8598b",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/4542 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Extracting ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw\n",
            "\n",
            "Rank 0 entering the 'run' function\n",
            "Rank 2 entering the 'run' function\n",
            "Rank 3 entering the 'run' function\n",
            "Rank 1 entering the 'run' function\n",
            "Computing validation set accuracy for epoch untrained\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 1875/1875 [00:19<00:00, 97.97it/s] \n",
            "100%|██████████| 313/313 [00:03<00:00, 98.35it/s] \n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Computing accuracy\n",
            "Validation set accuracy (Precision@1) = 0.9228\n",
            "Rank 3 starting epoch 0\n",
            "Rank 2 starting epoch 0\n",
            "Rank 0 starting epoch 0\n",
            "Rank 1 starting epoch 0\n",
            "Rank 3, iteration 0, loss 0.866588830947876, num pos pairs 6363, num neg pairs 58645\n",
            "Rank 0, iteration 0, loss 0.9615733027458191, num pos pairs 6779, num neg pairs 57892\n",
            "Rank 1, iteration 0, loss 0.8446443676948547, num pos pairs 6602, num neg pairs 58345\n",
            "Rank 2, iteration 0, loss 0.8583846092224121, num pos pairs 6573, num neg pairs 58385\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "INFO:root:Reducer buckets have been rebuilt in this iteration.\n",
            "INFO:root:Reducer buckets have been rebuilt in this iteration.\n",
            "INFO:root:Reducer buckets have been rebuilt in this iteration.\n",
            "INFO:root:Reducer buckets have been rebuilt in this iteration.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Rank 1, iteration 10, loss 0.3466784358024597, num pos pairs 6690, num neg pairs 53875\n",
            "Rank 3, iteration 10, loss 0.36289140582084656, num pos pairs 6817, num neg pairs 54370\n",
            "Rank 0, iteration 10, loss 0.3480638563632965, num pos pairs 6866, num neg pairs 53357\n",
            "Rank 2, iteration 10, loss 0.3387804925441742, num pos pairs 6831, num neg pairs 54233\n",
            "Rank 2, iteration 20, loss 0.39264407753944397, num pos pairs 6676, num neg pairs 43905\n",
            "Rank 0, iteration 20, loss 0.4453757703304291, num pos pairs 6646, num neg pairs 43683\n",
            "Rank 3, iteration 20, loss 0.42719975113868713, num pos pairs 6605, num neg pairs 45702\n",
            "Rank 1, iteration 20, loss 0.3941086232662201, num pos pairs 6669, num neg pairs 43935\n",
            "Rank 0, iteration 30, loss 0.3867945075035095, num pos pairs 6726, num neg pairs 43059\n",
            "Rank 1, iteration 30, loss 0.3979511260986328, num pos pairs 6801, num neg pairs 43861\n",
            "Rank 2, iteration 30, loss 0.40427878499031067, num pos pairs 6738, num neg pairs 43553\n",
            "Rank 3, iteration 30, loss 0.3638361394405365, num pos pairs 6729, num neg pairs 44446\n",
            "Rank 0, iteration 40, loss 0.40204280614852905, num pos pairs 6719, num neg pairs 32371\n",
            "Rank 1, iteration 40, loss 0.40012165904045105, num pos pairs 6756, num neg pairs 30024\n",
            "Rank 3, iteration 40, loss 0.4065723419189453, num pos pairs 6553, num neg pairs 32994\n",
            "Rank 2, iteration 40, loss 0.4115908741950989, num pos pairs 6679, num neg pairs 32125\n",
            "Rank 0, iteration 50, loss 0.4256925582885742, num pos pairs 6770, num neg pairs 32771\n",
            "Rank 1, iteration 50, loss 0.4157869517803192, num pos pairs 6424, num neg pairs 33140\n",
            "Rank 3, iteration 50, loss 0.4077575206756592, num pos pairs 6827, num neg pairs 31859\n",
            "Rank 2, iteration 50, loss 0.41882139444351196, num pos pairs 6578, num neg pairs 33105\n",
            "Rank 0, iteration 60, loss 0.3855322301387787, num pos pairs 6606, num neg pairs 28516\n",
            "Rank 1, iteration 60, loss 0.41611406207084656, num pos pairs 6675, num neg pairs 32194\n",
            "Rank 2, iteration 60, loss 0.3803099989891052, num pos pairs 6541, num neg pairs 33213\n",
            "Rank 3, iteration 60, loss 0.39375025033950806, num pos pairs 6551, num neg pairs 31639\n",
            "Rank 1, iteration 70, loss 0.41387301683425903, num pos pairs 6642, num neg pairs 29566\n",
            "Rank 0, iteration 70, loss 0.4156259596347809, num pos pairs 6512, num neg pairs 27427\n",
            "Rank 3, iteration 70, loss 0.4182702898979187, num pos pairs 6773, num neg pairs 29759\n",
            "Rank 2, iteration 70, loss 0.5334271192550659, num pos pairs 6512, num neg pairs 29644\n",
            "Rank 0, iteration 80, loss 0.4387570321559906, num pos pairs 6370, num neg pairs 27309\n",
            "Rank 2, iteration 80, loss 0.4312591552734375, num pos pairs 6368, num neg pairs 27025\n",
            "Rank 3, iteration 80, loss 0.4836696982383728, num pos pairs 6408, num neg pairs 27460\n",
            "Rank 1, iteration 80, loss 0.49333322048187256, num pos pairs 6450, num neg pairs 27822\n",
            "Rank 1, iteration 90, loss 0.3953418731689453, num pos pairs 6226, num neg pairs 26368\n",
            "Rank 2, iteration 90, loss 0.41428035497665405, num pos pairs 6279, num neg pairs 26909\n",
            "Rank 0, iteration 90, loss 0.46495315432548523, num pos pairs 6205, num neg pairs 25551\n",
            "Rank 3, iteration 90, loss 0.4484163522720337, num pos pairs 6367, num neg pairs 27692\n",
            "Rank 1, iteration 100, loss 0.42160147428512573, num pos pairs 5924, num neg pairs 18875\n",
            "Rank 3, iteration 100, loss 0.38765230774879456, num pos pairs 6174, num neg pairs 20279\n",
            "Rank 0, iteration 100, loss 0.452779084444046, num pos pairs 5483, num neg pairs 15033\n",
            "Rank 2, iteration 100, loss 0.4665995240211487, num pos pairs 5889, num neg pairs 20458\n",
            "Rank 0, iteration 110, loss 0.35404667258262634, num pos pairs 6344, num neg pairs 17951\n",
            "Rank 3, iteration 110, loss 0.40183010697364807, num pos pairs 6110, num neg pairs 18882\n",
            "Rank 2, iteration 110, loss 0.4514034688472748, num pos pairs 6119, num neg pairs 21353\n",
            "Rank 1, iteration 110, loss 0.43545687198638916, num pos pairs 6206, num neg pairs 19842\n",
            "Rank 1, epoch 0, average loss 0.4237597723633556\n",
            "Rank 3, epoch 0, average loss 0.4214024570281223\n",
            "Rank 0, epoch 0, average loss 0.4207431158777011\n",
            "Rank 2, epoch 0, average loss 0.418050453824512\n",
            "Computing validation set accuracy for epoch 0\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 1875/1875 [00:19<00:00, 96.98it/s]\n",
            "100%|██████████| 313/313 [00:03<00:00, 97.10it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Computing accuracy\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "INFO:PML:running k-nn with k=1\n",
            "INFO:PML:embedding dimensionality is 50\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Validation set accuracy (Precision@1) = 0.9775\n"
          ]
        }
      ],
      "source": [
        "import logging\n",
        "\n",
        "######################################################################################################################\n",
        "### This script is modified from the guide on pytorch distributed training https://github.com/seba-1511/dist_tuto.pth/\n",
        "### https://pytorch.org/tutorials/intermediate/dist_tuto.html\n",
        "######################################################################################################################\n",
        "import os\n",
        "from math import ceil\n",
        "from random import Random\n",
        "\n",
        "import torch\n",
        "import torch.distributed as dist\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.multiprocessing import Process\n",
        "from torch.nn.parallel import DistributedDataParallel as DDP\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "from pytorch_metric_learning import losses, miners, testers\n",
        "from pytorch_metric_learning.utils import distributed as pml_dist\n",
        "from pytorch_metric_learning.utils.accuracy_calculator import AccuracyCalculator\n",
        "\n",
        "logging.getLogger().setLevel(logging.INFO)\n",
        "\n",
        "\n",
        "class Partition(object):\n",
        "    \"\"\"Dataset-like object, but only access a subset of it.\"\"\"\n",
        "\n",
        "    def __init__(self, data, index):\n",
        "        self.data = data\n",
        "        self.index = index\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.index)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        data_idx = self.index[index]\n",
        "        return self.data[data_idx]\n",
        "\n",
        "\n",
        "class DataPartitioner(object):\n",
        "    \"\"\"Partitions a dataset into different chuncks.\"\"\"\n",
        "\n",
        "    def __init__(self, data, sizes=[0.7, 0.2, 0.1], seed=1234):\n",
        "        self.data = data\n",
        "        self.partitions = []\n",
        "        rng = Random()\n",
        "        rng.seed(seed)\n",
        "        data_len = len(data)\n",
        "        indexes = [x for x in range(0, data_len)]\n",
        "        rng.shuffle(indexes)\n",
        "\n",
        "        for frac in sizes:\n",
        "            part_len = int(frac * data_len)\n",
        "            self.partitions.append(indexes[0:part_len])\n",
        "            indexes = indexes[part_len:]\n",
        "\n",
        "    def use(self, partition):\n",
        "        return Partition(self.data, self.partitions[partition])\n",
        "\n",
        "\n",
        "class Net(nn.Module):\n",
        "    \"\"\"Network architecture.\"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
        "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
        "        self.conv2_drop = nn.Dropout2d()\n",
        "        self.fc1 = nn.Linear(320, 50)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
        "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
        "        x = x.view(-1, 320)\n",
        "        return self.fc1(x)\n",
        "\n",
        "\n",
        "def get_MNIST(train):\n",
        "    return datasets.MNIST(\n",
        "        \"./data\",\n",
        "        train=train,\n",
        "        download=True,\n",
        "        transform=transforms.Compose(\n",
        "            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
        "        ),\n",
        "    )\n",
        "\n",
        "\n",
        "def partition_dataset(dataset):\n",
        "    \"\"\"Partitioning MNIST\"\"\"\n",
        "    size = dist.get_world_size()\n",
        "    bsz = 512 // size\n",
        "    partition_sizes = [1.0 / size for _ in range(size)]\n",
        "    partition = DataPartitioner(dataset, partition_sizes)\n",
        "    partition = partition.use(dist.get_rank())\n",
        "    train_set = torch.utils.data.DataLoader(partition, batch_size=bsz, shuffle=True)\n",
        "    return train_set, bsz\n",
        "\n",
        "\n",
        "### convenient function from pytorch-metric-learning ###\n",
        "def get_all_embeddings(dataset, model, data_device):\n",
        "    # dataloader_num_workers has to be 0 to avoid pid error\n",
        "    # This only happens when within multiprocessing\n",
        "    tester = testers.BaseTester(dataloader_num_workers=0, data_device=data_device)\n",
        "    return tester.get_all_embeddings(dataset, model)\n",
        "\n",
        "\n",
        "### compute accuracy using AccuracyCalculator from pytorch-metric-learning ###\n",
        "def test(train_set, test_set, model, accuracy_calculator, data_device):\n",
        "    train_embeddings, train_labels = get_all_embeddings(train_set, model, data_device)\n",
        "    test_embeddings, test_labels = get_all_embeddings(test_set, model, data_device)\n",
        "    train_labels = train_labels.squeeze(1)\n",
        "    test_labels = test_labels.squeeze(1)\n",
        "    print(\"Computing accuracy\")\n",
        "    accuracies = accuracy_calculator.get_accuracy(\n",
        "        test_embeddings, test_labels, train_embeddings, train_labels, False\n",
        "    )\n",
        "    print(\n",
        "        \"Validation set accuracy (Precision@1) = {}\".format(\n",
        "            accuracies[\"precision_at_1\"]\n",
        "        )\n",
        "    )\n",
        "\n",
        "\n",
        "def test_model(rank, train_set, test_set, model, epoch, data_device):\n",
        "    if rank == 0:\n",
        "        print(\"Computing validation set accuracy for epoch {}\".format(epoch))\n",
        "        accuracy_calculator = AccuracyCalculator(include=(\"precision_at_1\",), k=1)\n",
        "        test(train_set, test_set, model, accuracy_calculator, data_device)\n",
        "    dist.barrier()\n",
        "\n",
        "\n",
        "def run(rank, size, train_dataset, val_dataset):\n",
        "    \"\"\"Distributed Synchronous SGD Example\"\"\"\n",
        "    print(\"Rank {} entering the 'run' function\".format(rank))\n",
        "    torch.manual_seed(1234)\n",
        "    train_set, bsz = partition_dataset(train_dataset)\n",
        "    dist.barrier()\n",
        "    ### use this if you have multiple GPUs ###\n",
        "    # device = torch.device(\"cuda:{}\".format(rank))\n",
        "    device = torch.device(\"cpu\")\n",
        "    model = Net()\n",
        "    ### if you have multiple GPUs, set this to DDP(model.to(device), device_ids=[rank])\n",
        "    model = DDP(model.to(device))\n",
        "    test_model(rank, train_dataset, val_dataset, model, \"untrained\", device)\n",
        "\n",
        "    optimizer = optim.Adam(model.parameters(), lr=0.01)\n",
        "\n",
        "    #####################################\n",
        "    ### pytorch-metric-learning stuff ###\n",
        "    loss_fn = losses.TripletMarginLoss()\n",
        "    loss_fn = pml_dist.DistributedLossWrapper(loss=loss_fn, efficient=True)\n",
        "    miner = miners.MultiSimilarityMiner()\n",
        "    miner = pml_dist.DistributedMinerWrapper(miner=miner, efficient=True)\n",
        "    ### pytorch-metric-learning stuff ###\n",
        "    #####################################\n",
        "\n",
        "    num_batches = ceil(len(train_set.dataset) / float(bsz))\n",
        "    for epoch in range(1):\n",
        "        epoch_loss = 0.0\n",
        "        print(\"Rank {} starting epoch {}\".format(rank, epoch))\n",
        "        for i, (data, target) in enumerate(train_set):\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            optimizer.zero_grad()\n",
        "            output = model(data)\n",
        "            hard_pairs = miner(output, target)\n",
        "            loss = loss_fn(output, target, hard_pairs)\n",
        "            epoch_loss += loss.item()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            if i % 10 == 0:\n",
        "                print(\n",
        "                    \"Rank {}, iteration {}, loss {}, num pos pairs {}, num neg pairs {}\".format(\n",
        "                        rank,\n",
        "                        i,\n",
        "                        loss.item(),\n",
        "                        miner.miner.num_pos_pairs,\n",
        "                        miner.miner.num_neg_pairs,\n",
        "                    )\n",
        "                )\n",
        "            dist.barrier()\n",
        "\n",
        "        print(\n",
        "            \"Rank {}, epoch {}, average loss {}\".format(\n",
        "                rank, epoch, epoch_loss / num_batches\n",
        "            )\n",
        "        )\n",
        "        test_model(rank, train_dataset, val_dataset, model, epoch, device)\n",
        "\n",
        "\n",
        "#######################################\n",
        "### Set backend='nccl' if using GPU ###\n",
        "#######################################\n",
        "def init_processes(rank, size, fn, train_dataset, val_dataset, backend=\"gloo\"):\n",
        "    \"\"\"Initialize the distributed environment.\"\"\"\n",
        "    os.environ[\"MASTER_ADDR\"] = \"localhost\"\n",
        "    os.environ[\"MASTER_PORT\"] = \"29500\"\n",
        "    dist.init_process_group(backend, rank=rank, world_size=size)\n",
        "    fn(rank, size, train_dataset, val_dataset)\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    train_dataset = get_MNIST(True)\n",
        "    val_dataset = get_MNIST(False)\n",
        "\n",
        "    size = 4\n",
        "    processes = []\n",
        "    for rank in range(size):\n",
        "        p = Process(\n",
        "            target=init_processes, args=(rank, size, run, train_dataset, val_dataset)\n",
        "        )\n",
        "        p.start()\n",
        "        processes.append(p)\n",
        "\n",
        "    for p in processes:\n",
        "        p.join()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "U0D96cL8co_D"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "DistributedTripletMarginLossMNIST.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.8.10"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "0107a7e435ff4f06a912b98e190c2c9d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_e626e3253ba1417596f6b5d3f87d8ebe",
              "IPY_MODEL_fc0e2d892d8c4779aaefc3660f53670e",
              "IPY_MODEL_6cb32fa30a4a4900a3d20b8ab3bdfe51"
            ],
            "layout": "IPY_MODEL_82a231510951443a882e0a4d7e942467"
          }
        },
        "0ddb2feda3c244439acda50c16d87732": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "22c6745c4d474207b36c4c3acea542ce": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "25409ddba9e54cea97b43de0e4a45947": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "257593fba66f47ef8ecf23b3b5528b07": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "26e7efdf822b42d99290b2b8f2d7c050": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_48d0cf3885884015b45c6c471d0e2b1b",
              "IPY_MODEL_aaa2a088e7de420f8bd6967033c87777",
              "IPY_MODEL_d6e739eee399488aaf1e3c627e1f3221"
            ],
            "layout": "IPY_MODEL_2a1a0715c79e473291feeecf52ede298"
          }
        },
        "2a1a0715c79e473291feeecf52ede298": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "30c1a16ac224445fadcc2d1f7f141936": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "339792cd03f2408b9310baae564127fe": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d38ad67181854718b18b0a88228743aa",
            "max": 4542,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_30c1a16ac224445fadcc2d1f7f141936",
            "value": 4542
          }
        },
        "38653e48529a475eb0756bccb5820ce6": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "39147f4e534040278af7f4be263da703": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "43de5b6f54dd4e69aaab95346425fa4d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "48d0cf3885884015b45c6c471d0e2b1b": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_9b861d5a4fa545a59ab089947c0f379e",
            "placeholder": "​",
            "style": "IPY_MODEL_257593fba66f47ef8ecf23b3b5528b07",
            "value": ""
          }
        },
        "48ea2d00bbaa46b7a52e4b3deed318ce": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "4a9be645d30545dc813dc71bc5eca1a3": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4c9ff1b060f34471a37d7ac5d19596fe": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_9fc0120ee9a94a09baedfeac38d0d9ff",
              "IPY_MODEL_8b528385a13e4e3aa91bf5fcffce3409",
              "IPY_MODEL_e372ec3e6cc44b5190241d47089d346c"
            ],
            "layout": "IPY_MODEL_ca98a6c54ef4473cb2b242d35ccc9fb6"
          }
        },
        "5d8e3ca6aec247528f8c67cd1efc7ffc": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6341bdeb007b449b9876cce5de240c92": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "6c4a3be721784447aa417eb861ed43af": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "6cb32fa30a4a4900a3d20b8ab3bdfe51": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_6fad8f3ff1884babbe4db769db6bdeae",
            "placeholder": "​",
            "style": "IPY_MODEL_96cda59489624e21b67628b44a184e8b",
            "value": " 29696/? [00:00&lt;00:00, 823617.66it/s]"
          }
        },
        "6fad8f3ff1884babbe4db769db6bdeae": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7574947398c443a9b0c71a59e5a8598b": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_a42ca39a12e44431a6ac7ae85a2e572d",
              "IPY_MODEL_339792cd03f2408b9310baae564127fe",
              "IPY_MODEL_91b6ab0c94194ca5a026170e703541b0"
            ],
            "layout": "IPY_MODEL_25409ddba9e54cea97b43de0e4a45947"
          }
        },
        "787746306a4248f6a995e215d68f2e2c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "82a231510951443a882e0a4d7e942467": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "8b528385a13e4e3aa91bf5fcffce3409": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a59365e2120c4b098c97c27742dc8de6",
            "max": 1648877,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_6c4a3be721784447aa417eb861ed43af",
            "value": 1648877
          }
        },
        "8dc29e4ab3bd4a6b8bf05270a206bb6e": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "91b6ab0c94194ca5a026170e703541b0": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8dc29e4ab3bd4a6b8bf05270a206bb6e",
            "placeholder": "​",
            "style": "IPY_MODEL_ce1056721ac04f1bb851a177501278eb",
            "value": " 5120/? [00:00&lt;00:00, 89515.04it/s]"
          }
        },
        "95e9f56ee5494094bb7780e426adc01d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "96cda59489624e21b67628b44a184e8b": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "9b861d5a4fa545a59ab089947c0f379e": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9fc0120ee9a94a09baedfeac38d0d9ff": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_38653e48529a475eb0756bccb5820ce6",
            "placeholder": "​",
            "style": "IPY_MODEL_22c6745c4d474207b36c4c3acea542ce",
            "value": ""
          }
        },
        "a42ca39a12e44431a6ac7ae85a2e572d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_da95ae822e624f0897d8c89ea77e3277",
            "placeholder": "​",
            "style": "IPY_MODEL_787746306a4248f6a995e215d68f2e2c",
            "value": ""
          }
        },
        "a59365e2120c4b098c97c27742dc8de6": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "aaa2a088e7de420f8bd6967033c87777": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4a9be645d30545dc813dc71bc5eca1a3",
            "max": 9912422,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_95e9f56ee5494094bb7780e426adc01d",
            "value": 9912422
          }
        },
        "c7f75365c3a54961a5698f0d8268e6a9": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ca98a6c54ef4473cb2b242d35ccc9fb6": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ce1056721ac04f1bb851a177501278eb": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d38ad67181854718b18b0a88228743aa": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d6e739eee399488aaf1e3c627e1f3221": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_39147f4e534040278af7f4be263da703",
            "placeholder": "​",
            "style": "IPY_MODEL_6341bdeb007b449b9876cce5de240c92",
            "value": " 9913344/? [00:00&lt;00:00, 22160775.86it/s]"
          }
        },
        "da95ae822e624f0897d8c89ea77e3277": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e372ec3e6cc44b5190241d47089d346c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5d8e3ca6aec247528f8c67cd1efc7ffc",
            "placeholder": "​",
            "style": "IPY_MODEL_43de5b6f54dd4e69aaab95346425fa4d",
            "value": " 1649664/? [00:00&lt;00:00, 7389854.87it/s]"
          }
        },
        "e626e3253ba1417596f6b5d3f87d8ebe": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c7f75365c3a54961a5698f0d8268e6a9",
            "placeholder": "​",
            "style": "IPY_MODEL_ea267ba5a5694e85a68dbdd76f84ad5d",
            "value": ""
          }
        },
        "ea267ba5a5694e85a68dbdd76f84ad5d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "fc0e2d892d8c4779aaefc3660f53670e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0ddb2feda3c244439acda50c16d87732",
            "max": 28881,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_48ea2d00bbaa46b7a52e4b3deed318ce",
            "value": 28881
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}