[skip ci] Updated assets/tldr/teaser.ipynb (#2871)

Author: vfdev-5

assets/tldr/teaser.ipynb

@@ -23,7 +23,7 @@
    "outputs": [],
    "source": [
     "# Install PyTorch-Ignite\n",
-    "!pip install pytorch-ignite"
+    "!pip install -q pytorch-ignite"
    ]
   },
   {
@@ -45,11 +45,13 @@
     "import os\n",
     "\n",
     "in_colab = \"COLAB_TPU_ADDR\" in os.environ\n",
-    "with_torch_launch = \"WORLD_SIZE\" in os.environ\n",
+    "with_torchrun = \"WORLD_SIZE\" in os.environ\n",
     "\n",
     "if in_colab:\n",
     "    VERSION = !curl -s https://api.github.com/repos/pytorch/xla/releases/latest | grep -Po '\"tag_name\": \"v\\K.*?(?=\")'\n",
-    "    !pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-{VERSION[0]}-cp37-cp37m-linux_x86_64.whl"
+    "    !pip install --upgrade -q cloud-tpu-client==0.10 torch=={VERSION[0]} torchvision https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{VERSION[0][:-2]}-cp38-cp38-linux_x86_64.whl\n",
+    "\n",
+    "!pip list | grep torch"
    ]
   },
   {
@@ -88,9 +90,18 @@
     "\n",
     "\n",
     "def get_train_test_datasets(path):\n",
+    "    # - Get train/test datasets\n",
+    "    if idist.get_rank() > 0:\n",
+    "        # Ensure that only rank 0 download the dataset\n",
+    "        idist.barrier()\n",
+    "\n",
     "    train_ds = datasets.CIFAR10(root=path, train=True, download=True, transform=train_transform)\n",
     "    test_ds = datasets.CIFAR10(root=path, train=False, download=False, transform=test_transform)\n",
     "\n",
+    "    if idist.get_rank() == 0:\n",
+    "        # Ensure that only rank 0 download the dataset\n",
+    "        idist.barrier()\n",
+    "\n",
     "    return train_ds, test_ds\n",
     "\n",
     "\n",
@@ -105,17 +116,8 @@
     "\n",
     "def get_dataflow(config):\n",
     "\n",
-    "    # - Get train/test datasets\n",
-    "    if idist.get_rank() > 0:\n",
-    "        # Ensure that only rank 0 download the dataset\n",
-    "        idist.barrier()\n",
-    "\n",
     "    train_dataset, test_dataset = get_train_test_datasets(config.get(\"data_path\", \".\"))\n",
     "\n",
-    "    if idist.get_rank() == 0:\n",
-    "        # Ensure that only rank 0 download the dataset\n",
-    "        idist.barrier()\n",
-    "\n",
     "    # Setup data loader also adapted to distributed config: nccl, gloo, xla-tpu\n",
     "    train_loader = idist.auto_dataloader(\n",
     "        train_dataset,\n",
@@ -402,10 +404,10 @@
     "# --- Single computation device ---\n",
     "# $ python main.py\n",
     "#\n",
-    "if __name__ == \"__main__\" and not (in_colab or with_torch_launch):\n",
+    "if __name__ == \"__main__\" and not (in_colab or with_torchrun):\n",
     "\n",
-    "    backend = None  # or \"nccl\", \"gloo\", \"xla-tpu\" ...\n",
-    "    nproc_per_node = None  # or N to spawn N processes\n",
+    "    backend = None\n",
+    "    nproc_per_node = None\n",
     "    config = {\n",
     "        \"model\": \"resnet18\",\n",
     "        \"dataset\": \"cifar10\",\n",
@@ -416,12 +418,12 @@
     "\n",
     "\n",
     "# --- Multiple GPUs ---\n",
-    "# $ python -m torch.distributed.launch --nproc_per_node=2 --use_env main.py\n",
+    "# $ torchrun --nproc_per_node=2 main.py\n",
     "#\n",
-    "if __name__ == \"__main__\" and with_torch_launch:\n",
+    "if __name__ == \"__main__\" and with_torchrun:\n",
     "\n",
-    "    backend = \"nccl\"  # or \"nccl\", \"gloo\", \"xla-tpu\" ...\n",
-    "    nproc_per_node = None  # or N to spawn N processes\n",
+    "    backend = \"nccl\"  # or \"nccl\", \"gloo\"\n",
+    "    nproc_per_node = None\n",
     "    config = {\n",
     "        \"model\": \"resnet18\",\n",
     "        \"dataset\": \"cifar10\",\n",
@@ -435,8 +437,8 @@
     "#\n",
     "if in_colab:\n",
     "\n",
-    "    backend = \"xla-tpu\"  # or \"nccl\", \"gloo\", \"xla-tpu\" ...\n",
-    "    nproc_per_node = 8  # or N to spawn N processes\n",
+    "    backend = \"xla-tpu\"\n",
+    "    nproc_per_node = 8\n",
     "    config = {\n",
     "        \"model\": \"resnet18\",\n",
     "        \"dataset\": \"cifar10\",\n",
@@ -465,7 +467,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3.10.6 64-bit",
    "language": "python",
    "name": "python3"
   },
@@ -479,7 +481,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.7"
+   "version": "3.10.6"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
+   }
   }
  },
  "nbformat": 4,

assets/tldr/teaser.py

@@ -15,7 +15,7 @@
 from ignite.metrics import Accuracy
 
 in_colab = "COLAB_TPU_ADDR" in os.environ
-with_torch_launch = "WORLD_SIZE" in os.environ
+with_torchrun = "WORLD_SIZE" in os.environ
 
 train_transform = Compose(
     [
@@ -31,9 +31,18 @@
 
 
 def get_train_test_datasets(path):
+    # - Get train/test datasets
+    if idist.get_rank() > 0:
+        # Ensure that only rank 0 download the dataset
+        idist.barrier()
+
     train_ds = datasets.CIFAR10(root=path, train=True, download=True, transform=train_transform)
     test_ds = datasets.CIFAR10(root=path, train=False, download=False, transform=test_transform)
 
+    if idist.get_rank() == 0:
+        # Ensure that only rank 0 download the dataset
+        idist.barrier()
+
     return train_ds, test_ds
 
 
@@ -48,17 +57,8 @@ def get_model(name):
 
 def get_dataflow(config):
 
-    # - Get train/test datasets
-    if idist.get_rank() > 0:
-        # Ensure that only rank 0 download the dataset
-        idist.barrier()
-
     train_dataset, test_dataset = get_train_test_datasets(config.get("data_path", "."))
 
-    if idist.get_rank() == 0:
-        # Ensure that only rank 0 download the dataset
-        idist.barrier()
-
     # Setup data loader also adapted to distributed config: nccl, gloo, xla-tpu
     train_loader = idist.auto_dataloader(
         train_dataset,
@@ -175,10 +175,10 @@ def evaluate_model():
 # --- Single computation device ---
 # $ python main.py
 #
-if __name__ == "__main__" and not (in_colab or with_torch_launch):
+if __name__ == "__main__" and not (in_colab or with_torchrun):
 
-    backend = None  # or "nccl", "gloo", "xla-tpu" ...
-    nproc_per_node = None  # or N to spawn N processes
+    backend = None
+    nproc_per_node = None
     config = {
         "model": "resnet18",
         "dataset": "cifar10",
@@ -189,12 +189,12 @@ def evaluate_model():
 
 
 # --- Multiple GPUs ---
-# $ python -m torch.distributed.launch --nproc_per_node=2 --use_env main.py
+# $ torchrun --nproc_per_node=2 main.py
 #
-if __name__ == "__main__" and with_torch_launch:
+if __name__ == "__main__" and with_torchrun:
 
-    backend = "nccl"  # or "nccl", "gloo", "xla-tpu" ...
-    nproc_per_node = None  # or N to spawn N processes
+    backend = "nccl"  # or "nccl", "gloo", ...
+    nproc_per_node = None
     config = {
         "model": "resnet18",
         "dataset": "cifar10",
@@ -208,8 +208,8 @@ def evaluate_model():
 #
 if in_colab:
 
-    backend = "xla-tpu"  # or "nccl", "gloo", "xla-tpu" ...
-    nproc_per_node = 8  # or N to spawn N processes
+    backend = "xla-tpu"
+    nproc_per_node = 8
     config = {
         "model": "resnet18",
         "dataset": "cifar10",