pytorch_DDP_training.ipynb

{
  "metadata": {
    "kernelspec": {
      "language": "python",
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.13",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "kaggle": {
      "accelerator": "nvidiaTeslaT4",
      "dataSources": [
        {
          "sourceId": 7749135,
          "sourceType": "datasetVersion",
          "datasetId": 4530293
        }
      ],
      "dockerImageVersionId": 30665,
      "isInternetEnabled": true,
      "language": "python",
      "sourceType": "notebook",
      "isGpuEnabled": true
    },
    "colab": {
      "name": "Pytorch DDP",
      "provenance": [],
      "include_colab_link": true
    }
  },
  "nbformat_minor": 0,
  "nbformat": 4,
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/DayDreamChaser/torch-action/blob/main/pytorch_DDP_training.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "source": [
        "\n",
        "# IMPORTANT: RUN THIS CELL IN ORDER TO IMPORT YOUR KAGGLE DATA SOURCES\n",
        "# TO THE CORRECT LOCATION (/kaggle/input) IN YOUR NOTEBOOK,\n",
        "# THEN FEEL FREE TO DELETE THIS CELL.\n",
        "# NOTE: THIS NOTEBOOK ENVIRONMENT DIFFERS FROM KAGGLE'S PYTHON\n",
        "# ENVIRONMENT SO THERE MAY BE MISSING LIBRARIES USED BY YOUR\n",
        "# NOTEBOOK.\n",
        "\n",
        "import os\n",
        "import sys\n",
        "from tempfile import NamedTemporaryFile\n",
        "from urllib.request import urlopen\n",
        "from urllib.parse import unquote, urlparse\n",
        "from urllib.error import HTTPError\n",
        "from zipfile import ZipFile\n",
        "import tarfile\n",
        "import shutil\n",
        "\n",
        "CHUNK_SIZE = 40960\n",
        "DATA_SOURCE_MAPPING = 'pydata:https%3A%2F%2Fstorage.googleapis.com%2Fkaggle-data-sets%2F4530293%2F7749135%2Fbundle%2Farchive.zip%3FX-Goog-Algorithm%3DGOOG4-RSA-SHA256%26X-Goog-Credential%3Dgcp-kaggle-com%2540kaggle-161607.iam.gserviceaccount.com%252F20240303%252Fauto%252Fstorage%252Fgoog4_request%26X-Goog-Date%3D20240303T092102Z%26X-Goog-Expires%3D259200%26X-Goog-SignedHeaders%3Dhost%26X-Goog-Signature%3D84f4378d65f7fc8d0dd430c421206cc61d8354316f0e551722866bec1a6db208de23539ad5305b50893112aab213e18a127f904bad0a4a1c138b4a19e72c25a6683ec9e4f142d3f70d5fed041f0c8c7d60419249402070fc53ac0132f8f3c1e4bb2e0c80ad41e95b29281fc750d03750bffbcbc66efa74bb229856eef66948c93b1132cd083836d809c5f3941d7dfe613f23de444f487c6177eb56a19ab97b2a56764c46ff38d742e455cb585b9e590590d66371df782cf7af993ee31bca94baad38483cc1c14504c97803c8ba15c96d242884814408a8e5b984ede807cefdcb058c04469cbb6ea79a218242ea891c3b68d349ceffd3f1202bd52f9a6bd6cb1c'\n",
        "\n",
        "KAGGLE_INPUT_PATH='/kaggle/input'\n",
        "KAGGLE_WORKING_PATH='/kaggle/working'\n",
        "KAGGLE_SYMLINK='kaggle'\n",
        "\n",
        "!umount /kaggle/input/ 2> /dev/null\n",
        "shutil.rmtree('/kaggle/input', ignore_errors=True)\n",
        "os.makedirs(KAGGLE_INPUT_PATH, 0o777, exist_ok=True)\n",
        "os.makedirs(KAGGLE_WORKING_PATH, 0o777, exist_ok=True)\n",
        "\n",
        "try:\n",
        "  os.symlink(KAGGLE_INPUT_PATH, os.path.join(\"..\", 'input'), target_is_directory=True)\n",
        "except FileExistsError:\n",
        "  pass\n",
        "try:\n",
        "  os.symlink(KAGGLE_WORKING_PATH, os.path.join(\"..\", 'working'), target_is_directory=True)\n",
        "except FileExistsError:\n",
        "  pass\n",
        "\n",
        "for data_source_mapping in DATA_SOURCE_MAPPING.split(','):\n",
        "    directory, download_url_encoded = data_source_mapping.split(':')\n",
        "    download_url = unquote(download_url_encoded)\n",
        "    filename = urlparse(download_url).path\n",
        "    destination_path = os.path.join(KAGGLE_INPUT_PATH, directory)\n",
        "    try:\n",
        "        with urlopen(download_url) as fileres, NamedTemporaryFile() as tfile:\n",
        "            total_length = fileres.headers['content-length']\n",
        "            print(f'Downloading {directory}, {total_length} bytes compressed')\n",
        "            dl = 0\n",
        "            data = fileres.read(CHUNK_SIZE)\n",
        "            while len(data) > 0:\n",
        "                dl += len(data)\n",
        "                tfile.write(data)\n",
        "                done = int(50 * dl / int(total_length))\n",
        "                sys.stdout.write(f\"\\r[{'=' * done}{' ' * (50-done)}] {dl} bytes downloaded\")\n",
        "                sys.stdout.flush()\n",
        "                data = fileres.read(CHUNK_SIZE)\n",
        "            if filename.endswith('.zip'):\n",
        "              with ZipFile(tfile) as zfile:\n",
        "                zfile.extractall(destination_path)\n",
        "            else:\n",
        "              with tarfile.open(tfile.name) as tarfile:\n",
        "                tarfile.extractall(destination_path)\n",
        "            print(f'\\nDownloaded and uncompressed: {directory}')\n",
        "    except HTTPError as e:\n",
        "        print(f'Failed to load (likely expired) {download_url} to path {destination_path}')\n",
        "        continue\n",
        "    except OSError as e:\n",
        "        print(f'Failed to load {download_url} to path {destination_path}')\n",
        "        continue\n",
        "\n",
        "print('Data source import complete.')\n"
      ],
      "metadata": {
        "id": "i8dW1IyFneTU"
      },
      "cell_type": "code",
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision import datasets\n",
        "from torchvision.transforms import ToTensor\n",
        "\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "print(f'Using {device} device')\n",
        "\n",
        "training_data = datasets.FashionMNIST(\n",
        "    root=\"data\",\n",
        "    train=True,\n",
        "    download=True,\n",
        "    transform=ToTensor()\n",
        ")\n",
        "\n",
        "test_data = datasets.FashionMNIST(\n",
        "    root=\"data\",\n",
        "    train=False,\n",
        "    download=True,\n",
        "    transform=ToTensor()\n",
        ")\n",
        "\n",
        "learning_rate = 1e-3\n",
        "batch_size = 64\n",
        "epochs = 3\n",
        "\n",
        "train_dataloader = DataLoader(training_data, batch_size=batch_size)\n",
        "test_dataloader = DataLoader(test_data, batch_size=batch_size)"
      ],
      "metadata": {
        "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
        "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
        "execution": {
          "iopub.status.busy": "2024-03-03T08:44:33.642334Z",
          "iopub.execute_input": "2024-03-03T08:44:33.643149Z",
          "iopub.status.idle": "2024-03-03T08:44:38.881817Z",
          "shell.execute_reply.started": "2024-03-03T08:44:33.643099Z",
          "shell.execute_reply": "2024-03-03T08:44:38.880844Z"
        },
        "trusted": true,
        "id": "SIzW9ocQneTX",
        "outputId": "2d5b7d46-ff16-41f9-ddd2-1e0828563515"
      },
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "text": "Using cuda device\nDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz\nDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to data/FashionMNIST/raw/train-images-idx3-ubyte.gz\n",
          "output_type": "stream"
        },
        {
          "name": "stderr",
          "text": "100%|██████████| 26421880/26421880 [00:01<00:00, 17077834.30it/s]\n",
          "output_type": "stream"
        },
        {
          "name": "stdout",
          "text": "Extracting data/FashionMNIST/raw/train-images-idx3-ubyte.gz to data/FashionMNIST/raw\n\nDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz\nDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw/train-labels-idx1-ubyte.gz\n",
          "output_type": "stream"
        },
        {
          "name": "stderr",
          "text": "100%|██████████| 29515/29515 [00:00<00:00, 271325.90it/s]\n",
          "output_type": "stream"
        },
        {
          "name": "stdout",
          "text": "Extracting data/FashionMNIST/raw/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw\n\nDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz\nDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz\n",
          "output_type": "stream"
        },
        {
          "name": "stderr",
          "text": "100%|██████████| 4422102/4422102 [00:00<00:00, 4994955.96it/s]\n",
          "output_type": "stream"
        },
        {
          "name": "stdout",
          "text": "Extracting data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw\n\nDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz\nDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz\n",
          "output_type": "stream"
        },
        {
          "name": "stderr",
          "text": "100%|██████████| 5148/5148 [00:00<00:00, 10678673.09it/s]",
          "output_type": "stream"
        },
        {
          "name": "stdout",
          "text": "Extracting data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw\n\n",
          "output_type": "stream"
        },
        {
          "name": "stderr",
          "text": "\n",
          "output_type": "stream"
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class NeuralNetwork(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(NeuralNetwork, self).__init__()\n",
        "        self.flatten = nn.Flatten()\n",
        "        self.linear_relu_stack = nn.Sequential(\n",
        "            nn.Linear(28*28, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, 256),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(256, 10),\n",
        "            nn.Dropout(p=0.2)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.flatten(x)\n",
        "        logits = self.linear_relu_stack(x)\n",
        "        return logits\n",
        "\n",
        "model = NeuralNetwork().to(device)\n",
        "\n",
        "def train_loop(dataloader, model, loss_fn, optimizer):\n",
        "    size = len(dataloader.dataset)\n",
        "    model.train()\n",
        "    for batch, (X, y) in enumerate(dataloader, start=1):\n",
        "        X, y = X.to(device), y.to(device)\n",
        "        # Compute prediction and loss\n",
        "        pred = model(X)\n",
        "        loss = loss_fn(pred, y)\n",
        "        # Backpropagation\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        if batch % 100 == 0:\n",
        "            loss, current = loss.item(), batch * len(X)\n",
        "            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")\n",
        "\n",
        "def test_loop(dataloader, model, loss_fn):\n",
        "    size = len(dataloader.dataset)\n",
        "    num_batches = len(dataloader)\n",
        "    test_loss, correct = 0, 0\n",
        "\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "        for X, y in dataloader:\n",
        "            X, y = X.to(device), y.to(device)\n",
        "            pred = model(X)\n",
        "            test_loss += loss_fn(pred, y).item()\n",
        "            correct += (pred.argmax(dim=-1) == y).type(torch.float).sum().item()\n",
        "\n",
        "    test_loss /= num_batches\n",
        "    correct /= size\n",
        "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")\n",
        "\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
        "\n",
        "for t in range(epochs):\n",
        "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
        "    train_loop(train_dataloader, model, loss_fn, optimizer)\n",
        "    test_loop(test_dataloader, model, loss_fn)k\n",
        "print(\"Done!\")"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-03T08:44:45.397218Z",
          "iopub.execute_input": "2024-03-03T08:44:45.398001Z",
          "iopub.status.idle": "2024-03-03T08:45:14.367899Z",
          "shell.execute_reply.started": "2024-03-03T08:44:45.397961Z",
          "shell.execute_reply": "2024-03-03T08:45:14.367052Z"
        },
        "trusted": true,
        "id": "EAfZYIgpneTY",
        "outputId": "7ab0789e-0eb3-4892-ba25-6aab8d19b1cf"
      },
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "text": "Epoch 1\n-------------------------------\nloss: 0.992364  [ 6400/60000]\nloss: 0.897658  [12800/60000]\nloss: 0.590443  [19200/60000]\nloss: 0.959099  [25600/60000]\nloss: 0.425988  [32000/60000]\nloss: 0.526136  [38400/60000]\nloss: 0.861878  [44800/60000]\nloss: 0.628830  [51200/60000]\nloss: 0.629240  [57600/60000]\nTest Error: \n Accuracy: 83.6%, Avg loss: 0.454017 \n\nEpoch 2\n-------------------------------\nloss: 0.581930  [ 6400/60000]\nloss: 0.594764  [12800/60000]\nloss: 0.632375  [19200/60000]\nloss: 0.781220  [25600/60000]\nloss: 0.426663  [32000/60000]\nloss: 0.464518  [38400/60000]\nloss: 0.947383  [44800/60000]\nloss: 0.528568  [51200/60000]\nloss: 0.494135  [57600/60000]\nTest Error: \n Accuracy: 85.2%, Avg loss: 0.410554 \n\nEpoch 3\n-------------------------------\nloss: 0.610657  [ 6400/60000]\nloss: 0.697756  [12800/60000]\nloss: 0.468149  [19200/60000]\nloss: 0.720574  [25600/60000]\nloss: 0.367921  [32000/60000]\nloss: 0.493198  [38400/60000]\nloss: 0.697919  [44800/60000]\nloss: 0.519677  [51200/60000]\nloss: 0.463810  [57600/60000]\nTest Error: \n Accuracy: 85.9%, Avg loss: 0.388260 \n\nDone!\n",
          "output_type": "stream"
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-03T08:47:21.715963Z",
          "iopub.execute_input": "2024-03-03T08:47:21.716801Z",
          "iopub.status.idle": "2024-03-03T08:47:21.722263Z",
          "shell.execute_reply.started": "2024-03-03T08:47:21.716771Z",
          "shell.execute_reply": "2024-03-03T08:47:21.721249Z"
        },
        "trusted": true,
        "id": "D6eBoHBFneTZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os,PIL\n",
        "import numpy as np\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "import torchvision\n",
        "from torchvision import transforms\n",
        "import datetime\n",
        "\n",
        "#======================================================================\n",
        "# import accelerate\n",
        "from accelerate import Accelerator\n",
        "from accelerate.utils import set_seed\n",
        "#======================================================================\n",
        "\n",
        "\n",
        "def create_dataloaders(batch_size=64):\n",
        "    transform = transforms.Compose([transforms.ToTensor()])\n",
        "\n",
        "    ds_train = torchvision.datasets.MNIST(root=\"./minist/\",train=True,download=True,transform=transform)\n",
        "    ds_val = torchvision.datasets.MNIST(root=\"./minist/\",train=False,download=True,transform=transform)\n",
        "\n",
        "    dl_train =  torch.utils.data.DataLoader(ds_train, batch_size=batch_size, shuffle=True,\n",
        "                                            num_workers=2,drop_last=True)\n",
        "    dl_val =  torch.utils.data.DataLoader(ds_val, batch_size=batch_size, shuffle=False,\n",
        "                                          num_workers=2,drop_last=True)\n",
        "    return dl_train,dl_val\n",
        "\n",
        "\n",
        "def create_net():\n",
        "    net = nn.Sequential()\n",
        "    net.add_module(\"conv1\",nn.Conv2d(in_channels=1,out_channels=512,kernel_size = 3))\n",
        "    net.add_module(\"pool1\",nn.MaxPool2d(kernel_size = 2,stride = 2))\n",
        "    net.add_module(\"conv2\",nn.Conv2d(in_channels=512,out_channels=256,kernel_size = 5))\n",
        "    net.add_module(\"pool2\",nn.MaxPool2d(kernel_size = 2,stride = 2))\n",
        "    net.add_module(\"dropout\",nn.Dropout2d(p = 0.1))\n",
        "    net.add_module(\"adaptive_pool\",nn.AdaptiveMaxPool2d((1,1)))\n",
        "    net.add_module(\"flatten\",nn.Flatten())\n",
        "    net.add_module(\"linear1\",nn.Linear(256,128))\n",
        "    net.add_module(\"relu\",nn.ReLU())\n",
        "    net.add_module(\"linear2\",nn.Linear(128,10))\n",
        "    return net\n",
        "\n",
        "\n",
        "\n",
        "def training_loop(epochs = 5,\n",
        "                  lr = 1e-3,\n",
        "                  batch_size= 1024,\n",
        "                  ckpt_path = \"checkpoint.pt\",\n",
        "                  mixed_precision=\"no\", #'fp16'\n",
        "                 ):\n",
        "\n",
        "    train_dataloader, eval_dataloader = create_dataloaders(batch_size)\n",
        "    model = create_net()\n",
        "\n",
        "\n",
        "    optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)\n",
        "    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer=optimizer, max_lr=25*lr,\n",
        "                              epochs=epochs, steps_per_epoch=len(train_dataloader))\n",
        "\n",
        "    #======================================================================\n",
        "    # initialize accelerator and auto move data/model to accelerator.device\n",
        "    set_seed(42)\n",
        "    accelerator = Accelerator(mixed_precision=mixed_precision)\n",
        "    accelerator.print(f'device {str(accelerator.device)} is used!')\n",
        "    model, optimizer,lr_scheduler, train_dataloader, eval_dataloader = accelerator.prepare(\n",
        "        model, optimizer,lr_scheduler, train_dataloader, eval_dataloader)\n",
        "    #======================================================================\n",
        "\n",
        "\n",
        "    for epoch in range(epochs):\n",
        "        model.train()\n",
        "        for step, batch in enumerate(train_dataloader):\n",
        "            features,labels = batch\n",
        "            preds = model(features)\n",
        "            loss = nn.CrossEntropyLoss()(preds,labels)\n",
        "\n",
        "            #======================================================================\n",
        "            #attention here!\n",
        "            accelerator.backward(loss) #loss.backward()\n",
        "            #======================================================================\n",
        "\n",
        "            optimizer.step()\n",
        "            lr_scheduler.step()\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "\n",
        "\n",
        "        model.eval()\n",
        "        accurate = 0\n",
        "        num_elems = 0\n",
        "\n",
        "        for _, batch in enumerate(eval_dataloader):\n",
        "            features,labels = batch\n",
        "            with torch.no_grad():\n",
        "                preds = model(features)\n",
        "            predictions = preds.argmax(dim=-1)\n",
        "\n",
        "            #======================================================================\n",
        "            #gather data from multi-gpus (used when in ddp mode)\n",
        "            predictions = accelerator.gather(predictions)\n",
        "            labels = accelerator.gather(labels)\n",
        "            #======================================================================\n",
        "\n",
        "            accurate_preds =  (predictions==labels)\n",
        "            num_elems += accurate_preds.shape[0]\n",
        "            accurate += accurate_preds.long().sum()\n",
        "\n",
        "        eval_metric = accurate.item() / num_elems\n",
        "\n",
        "        #======================================================================\n",
        "        #print logs and save ckpt\n",
        "        accelerator.wait_for_everyone()\n",
        "        nowtime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n",
        "        accelerator.print(f\"epoch【{epoch}】@{nowtime} --> eval_metric= {100 * eval_metric:.2f}%\")\n",
        "        unwrapped_net = accelerator.unwrap_model(model)\n",
        "        accelerator.save(unwrapped_net.state_dict(),ckpt_path+\"_\"+str(epoch))\n",
        "        #======================================================================\n",
        "\n",
        "training_loop(epochs = 5,lr = 1e-3,batch_size= 1024,ckpt_path = \"checkpoint.pt\",\n",
        "            mixed_precision=\"no\") #mixed_precision='fp16'\n"
      ],
      "metadata": {
        "trusted": true,
        "id": "3yB4gL-wneTZ",
        "outputId": "d4a6c47d-56f3-4685-aac6-6eab04c053f7"
      },
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "text": "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\nDownloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./minist/MNIST/raw/train-images-idx3-ubyte.gz\n",
          "output_type": "stream"
        },
        {
          "name": "stderr",
          "text": "100%|██████████| 9912422/9912422 [00:00<00:00, 115489370.00it/s]\n",
          "output_type": "stream"
        },
        {
          "name": "stdout",
          "text": "Extracting ./minist/MNIST/raw/train-images-idx3-ubyte.gz to ./minist/MNIST/raw\n\nDownloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\nDownloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./minist/MNIST/raw/train-labels-idx1-ubyte.gz\n",
          "output_type": "stream"
        },
        {
          "name": "stderr",
          "text": "100%|██████████| 28881/28881 [00:00<00:00, 45149345.44it/s]",
          "output_type": "stream"
        },
        {
          "name": "stdout",
          "text": "Extracting ./minist/MNIST/raw/train-labels-idx1-ubyte.gz to ./minist/MNIST/raw\n\nDownloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
          "output_type": "stream"
        },
        {
          "name": "stderr",
          "text": "\n",
          "output_type": "stream"
        },
        {
          "name": "stdout",
          "text": "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./minist/MNIST/raw/t10k-images-idx3-ubyte.gz\n",
          "output_type": "stream"
        },
        {
          "name": "stderr",
          "text": "100%|██████████| 1648877/1648877 [00:00<00:00, 28517728.59it/s]\n",
          "output_type": "stream"
        },
        {
          "name": "stdout",
          "text": "Extracting ./minist/MNIST/raw/t10k-images-idx3-ubyte.gz to ./minist/MNIST/raw\n\nDownloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\nDownloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./minist/MNIST/raw/t10k-labels-idx1-ubyte.gz\n",
          "output_type": "stream"
        },
        {
          "name": "stderr",
          "text": "100%|██████████| 4542/4542 [00:00<00:00, 8172685.01it/s]\n",
          "output_type": "stream"
        },
        {
          "name": "stdout",
          "text": "Extracting ./minist/MNIST/raw/t10k-labels-idx1-ubyte.gz to ./minist/MNIST/raw\n\ndevice cuda is used!\nepoch【0】@2024-03-03 09:08:10 --> eval_metric= 10.18%\nepoch【1】@2024-03-03 09:08:24 --> eval_metric= 11.30%\nepoch【2】@2024-03-03 09:08:38 --> eval_metric= 11.30%\nepoch【3】@2024-03-03 09:08:52 --> eval_metric= 11.30%\nepoch【4】@2024-03-03 09:09:06 --> eval_metric= 11.30%\n",
          "output_type": "stream"
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "net = create_net()\n",
        "print(net)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-03T09:15:52.014416Z",
          "iopub.execute_input": "2024-03-03T09:15:52.01543Z",
          "iopub.status.idle": "2024-03-03T09:15:52.060657Z",
          "shell.execute_reply.started": "2024-03-03T09:15:52.015386Z",
          "shell.execute_reply": "2024-03-03T09:15:52.059646Z"
        },
        "trusted": true,
        "id": "uAOUQYRWneTa",
        "outputId": "24a15e43-5dba-4d3b-d5e8-255c151095b2"
      },
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "text": "Sequential(\n  (conv1): Conv2d(1, 512, kernel_size=(3, 3), stride=(1, 1))\n  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n  (conv2): Conv2d(512, 256, kernel_size=(5, 5), stride=(1, 1))\n  (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n  (dropout): Dropout2d(p=0.1, inplace=False)\n  (adaptive_pool): AdaptiveMaxPool2d(output_size=(1, 1))\n  (flatten): Flatten(start_dim=1, end_dim=-1)\n  (linear1): Linear(in_features=256, out_features=128, bias=True)\n  (relu): ReLU()\n  (linear2): Linear(in_features=128, out_features=10, bias=True)\n)\n",
          "output_type": "stream"
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "dl_train,dl_val = create_dataloaders(batch_size=64)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-03T09:16:22.00052Z",
          "iopub.execute_input": "2024-03-03T09:16:22.001035Z",
          "iopub.status.idle": "2024-03-03T09:16:22.092278Z",
          "shell.execute_reply.started": "2024-03-03T09:16:22.001002Z",
          "shell.execute_reply": "2024-03-03T09:16:22.091318Z"
        },
        "trusted": true,
        "id": "VSO389AqneTa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(dl_train.dataset[0])"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-03T09:19:54.050176Z",
          "iopub.execute_input": "2024-03-03T09:19:54.050564Z",
          "iopub.status.idle": "2024-03-03T09:19:54.094122Z",
          "shell.execute_reply.started": "2024-03-03T09:19:54.050537Z",
          "shell.execute_reply": "2024-03-03T09:19:54.093041Z"
        },
        "trusted": true,
        "id": "1D0MusxAneTb",
        "outputId": "a6c44891-c4f3-40c5-f20c-f11ba8a23085"
      },
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "text": "(tensor([[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0118, 0.0706, 0.0706, 0.0706,\n          0.4941, 0.5333, 0.6863, 0.1020, 0.6510, 1.0000, 0.9686, 0.4980,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.1176, 0.1412, 0.3686, 0.6039, 0.6667, 0.9922, 0.9922, 0.9922,\n          0.9922, 0.9922, 0.8824, 0.6745, 0.9922, 0.9490, 0.7647, 0.2510,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1922,\n          0.9333, 0.9922, 0.9922, 0.9922, 0.9922, 0.9922, 0.9922, 0.9922,\n          0.9922, 0.9843, 0.3647, 0.3216, 0.3216, 0.2196, 0.1529, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0706,\n          0.8588, 0.9922, 0.9922, 0.9922, 0.9922, 0.9922, 0.7765, 0.7137,\n          0.9686, 0.9451, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.3137, 0.6118, 0.4196, 0.9922, 0.9922, 0.8039, 0.0431, 0.0000,\n          0.1686, 0.6039, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0549, 0.0039, 0.6039, 0.9922, 0.3529, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.5451, 0.9922, 0.7451, 0.0078, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0431, 0.7451, 0.9922, 0.2745, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.1373, 0.9451, 0.8824, 0.6275,\n          0.4235, 0.0039, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.3176, 0.9412, 0.9922,\n          0.9922, 0.4667, 0.0980, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1765, 0.7294,\n          0.9922, 0.9922, 0.5882, 0.1059, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0627,\n          0.3647, 0.9882, 0.9922, 0.7333, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.9765, 0.9922, 0.9765, 0.2510, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1804, 0.5098,\n          0.7176, 0.9922, 0.9922, 0.8118, 0.0078, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.1529, 0.5804, 0.8980, 0.9922,\n          0.9922, 0.9922, 0.9804, 0.7137, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0941, 0.4471, 0.8667, 0.9922, 0.9922, 0.9922,\n          0.9922, 0.7882, 0.3059, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0902, 0.2588, 0.8353, 0.9922, 0.9922, 0.9922, 0.9922, 0.7765,\n          0.3176, 0.0078, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0706, 0.6706,\n          0.8588, 0.9922, 0.9922, 0.9922, 0.9922, 0.7647, 0.3137, 0.0353,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.2157, 0.6745, 0.8863, 0.9922,\n          0.9922, 0.9922, 0.9922, 0.9569, 0.5216, 0.0431, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.5333, 0.9922, 0.9922, 0.9922,\n          0.8314, 0.5294, 0.5176, 0.0627, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000],\n         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,\n          0.0000, 0.0000, 0.0000, 0.0000]]]), 5)\n",
          "output_type": "stream"
        },
        {
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[11], line 5\u001b[0m\n\u001b[1;32m      3\u001b[0m image \u001b[38;5;241m=\u001b[39m dl_train\u001b[38;5;241m.\u001b[39mdataset[\u001b[38;5;241m0\u001b[39m][\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mcpu()\u001b[38;5;241m.\u001b[39mclone()  \u001b[38;5;66;03m# we clone the tensor to not do changes on it\u001b[39;00m\n\u001b[1;32m      4\u001b[0m image \u001b[38;5;241m=\u001b[39m image\u001b[38;5;241m.\u001b[39msqueeze(\u001b[38;5;241m0\u001b[39m)  \u001b[38;5;66;03m# remove the fake batch dimension\u001b[39;00m\n\u001b[0;32m----> 5\u001b[0m image \u001b[38;5;241m=\u001b[39m \u001b[43munloader\u001b[49m(image)\n\u001b[1;32m      6\u001b[0m plt\u001b[38;5;241m.\u001b[39mimshow(image)\n\u001b[1;32m      8\u001b[0m \u001b[38;5;66;03m# ds_train = torchvision.datasets.MNIST(root=\"./minist/\",train=True,download=True,transform=transform)\u001b[39;00m\n",
            "\u001b[0;31mNameError\u001b[0m: name 'unloader' is not defined"
          ],
          "ename": "NameError",
          "evalue": "name 'unloader' is not defined",
          "output_type": "error"
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from accelerate.utils import write_basic_config\n",
        "write_basic_config() # Write a config file\n",
        "os._exit(0) # Restart the notebook to reload info from the latest config file\n"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-03T09:10:20.30641Z",
          "iopub.execute_input": "2024-03-03T09:10:20.307152Z"
        },
        "trusted": true,
        "id": "Lq_zy9TuneTb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# or answer some question to create a config\n",
        "#!accelerate config"
      ],
      "metadata": {
        "id": "LlxJjmKtneTb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os,PIL\n",
        "import numpy as np\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "import torchvision\n",
        "from torchvision import transforms\n",
        "import datetime\n",
        "\n",
        "#======================================================================\n",
        "# import accelerate\n",
        "from accelerate import Accelerator\n",
        "from accelerate.utils import set_seed\n",
        "#======================================================================\n",
        "\n",
        "\n",
        "def create_dataloaders(batch_size=64):\n",
        "    transform = transforms.Compose([transforms.ToTensor()])\n",
        "\n",
        "    ds_train = torchvision.datasets.MNIST(root=\"./minist/\",train=True,download=True,transform=transform)\n",
        "    ds_val = torchvision.datasets.MNIST(root=\"./minist/\",train=False,download=True,transform=transform)\n",
        "\n",
        "    dl_train =  torch.utils.data.DataLoader(ds_train, batch_size=batch_size, shuffle=True,\n",
        "                                            num_workers=2,drop_last=True)\n",
        "    dl_val =  torch.utils.data.DataLoader(ds_val, batch_size=batch_size, shuffle=False,\n",
        "                                          num_workers=2,drop_last=True)\n",
        "    return dl_train,dl_val\n",
        "\n",
        "\n",
        "def create_net():\n",
        "    net = nn.Sequential()\n",
        "    net.add_module(\"conv1\",nn.Conv2d(in_channels=1,out_channels=512,kernel_size = 3))\n",
        "    net.add_module(\"pool1\",nn.MaxPool2d(kernel_size = 2,stride = 2))\n",
        "    net.add_module(\"conv2\",nn.Conv2d(in_channels=512,out_channels=256,kernel_size = 5))\n",
        "    net.add_module(\"pool2\",nn.MaxPool2d(kernel_size = 2,stride = 2))\n",
        "    net.add_module(\"dropout\",nn.Dropout2d(p = 0.1))\n",
        "    net.add_module(\"adaptive_pool\",nn.AdaptiveMaxPool2d((1,1)))\n",
        "    net.add_module(\"flatten\",nn.Flatten())\n",
        "    net.add_module(\"linear1\",nn.Linear(256,128))\n",
        "    net.add_module(\"relu\",nn.ReLU())\n",
        "    net.add_module(\"linear2\",nn.Linear(128,10))\n",
        "    return net\n",
        "\n",
        "\n",
        "\n",
        "def training_loop(epochs = 5,\n",
        "                  lr = 1e-3,\n",
        "                  batch_size= 1024,\n",
        "                  ckpt_path = \"checkpoint.pt\",\n",
        "                  mixed_precision=\"no\", #'fp16'\n",
        "                 ):\n",
        "\n",
        "    train_dataloader, eval_dataloader = create_dataloaders(batch_size)\n",
        "    model = create_net()\n",
        "\n",
        "\n",
        "    optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)\n",
        "    lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer=optimizer, max_lr=25*lr,\n",
        "                              epochs=epochs, steps_per_epoch=len(train_dataloader))\n",
        "\n",
        "    #======================================================================\n",
        "    # initialize accelerator and auto move data/model to accelerator.device\n",
        "    set_seed(42)\n",
        "    accelerator = Accelerator(mixed_precision=mixed_precision)\n",
        "    accelerator.print(f'device {str(accelerator.device)} is used!')\n",
        "    model, optimizer,lr_scheduler, train_dataloader, eval_dataloader = accelerator.prepare(\n",
        "        model, optimizer,lr_scheduler, train_dataloader, eval_dataloader)\n",
        "    #======================================================================\n",
        "\n",
        "\n",
        "    for epoch in range(epochs):\n",
        "        model.train()\n",
        "        for step, batch in enumerate(train_dataloader):\n",
        "            features,labels = batch\n",
        "            preds = model(features)\n",
        "            loss = nn.CrossEntropyLoss()(preds,labels)\n",
        "\n",
        "            #======================================================================\n",
        "            #attention here!\n",
        "            accelerator.backward(loss) #loss.backward()\n",
        "            #======================================================================\n",
        "\n",
        "            optimizer.step()\n",
        "            lr_scheduler.step()\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "\n",
        "\n",
        "        model.eval()\n",
        "        accurate = 0\n",
        "        num_elems = 0\n",
        "\n",
        "        for _, batch in enumerate(eval_dataloader):\n",
        "            features,labels = batch\n",
        "            with torch.no_grad():\n",
        "                preds = model(features)\n",
        "            predictions = preds.argmax(dim=-1)\n",
        "\n",
        "            #======================================================================\n",
        "            #gather data from multi-gpus (used when in ddp mode)\n",
        "            predictions = accelerator.gather(predictions)\n",
        "            labels = accelerator.gather(labels)\n",
        "            #======================================================================\n",
        "\n",
        "            accurate_preds =  (predictions==labels)\n",
        "            num_elems += accurate_preds.shape[0]\n",
        "            accurate += accurate_preds.long().sum()\n",
        "\n",
        "        eval_metric = accurate.item() / num_elems\n",
        "\n",
        "        #======================================================================\n",
        "        #print logs and save ckpt\n",
        "        accelerator.wait_for_everyone()\n",
        "        nowtime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')\n",
        "        accelerator.print(f\"epoch【{epoch}】@{nowtime} --> eval_metric= {100 * eval_metric:.2f}%\")\n",
        "        unwrapped_net = accelerator.unwrap_model(model)\n",
        "        accelerator.save(unwrapped_net.state_dict(),ckpt_path+\"_\"+str(epoch))\n",
        "        #======================================================================\n",
        "\n",
        "# training_loop(epochs = 5,lr = 1e-4,batch_size= 1024,ckpt_path = \"checkpoint.pt\",\n",
        "#             mixed_precision=\"no\") #mixed_precision='fp16'"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-03T09:11:20.097703Z",
          "iopub.execute_input": "2024-03-03T09:11:20.098133Z",
          "iopub.status.idle": "2024-03-03T09:11:23.970535Z",
          "shell.execute_reply.started": "2024-03-03T09:11:20.098102Z",
          "shell.execute_reply": "2024-03-03T09:11:23.969732Z"
        },
        "trusted": true,
        "id": "2FZYlb_sneTc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from accelerate import notebook_launcher\n",
        "#args = (5,1e-4,1024,'checkpoint.pt','no')\n",
        "\n",
        "args = dict(epochs = 5,\n",
        "        lr = 1e-4,\n",
        "        batch_size= 1024,\n",
        "        ckpt_path = \"checkpoint.pt\",\n",
        "        mixed_precision=\"no\").values()\n",
        "notebook_launcher(training_loop, args, num_processes=2)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-03T09:11:31.276201Z",
          "iopub.execute_input": "2024-03-03T09:11:31.276696Z",
          "iopub.status.idle": "2024-03-03T09:12:12.967572Z",
          "shell.execute_reply.started": "2024-03-03T09:11:31.276668Z",
          "shell.execute_reply": "2024-03-03T09:12:12.96639Z"
        },
        "trusted": true,
        "id": "5XHK-EZzneTc",
        "outputId": "8737de9e-66ee-4831-f865-098fe1595caf"
      },
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "text": "Launching training on 2 GPUs.\ndevice cuda:0 is used!\nepoch【0】@2024-03-03 09:11:41 --> eval_metric= 90.37%\nepoch【1】@2024-03-03 09:11:49 --> eval_metric= 97.27%\nepoch【2】@2024-03-03 09:11:56 --> eval_metric= 98.10%\nepoch【3】@2024-03-03 09:12:04 --> eval_metric= 98.33%\nepoch【4】@2024-03-03 09:12:12 --> eval_metric= 98.43%\n",
          "output_type": "stream"
        }
      ]
    }
  ]
}