DPIS_Opacus v1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import libraries\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torch.autograd import Variable\n",
    "import time\n",
    "import torch.optim as optim\n",
    "import os\n",
    "from opacus import PrivacyEngine\n",
    "import math\n",
    "import random\n",
    "from opacus import GradSampleModule\n",
    "from opacus.validators import ModuleValidator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Device configuration\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model\n",
    "def conv3x3(in_channels, out_channels, stride=1):\n",
    "    \"\"\"\n",
    "    return 3x3 Conv2d\n",
    "    \"\"\"\n",
    "    return nn.Conv2d(in_channels, out_channels, kernel_size=3,stride=stride, padding=1, bias=False)\n",
    "\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    \"\"\"\n",
    "    Initialize basic ResidualBlock with forward propogation\n",
    "    \"\"\"\n",
    "    def __init__(self, in_channels, out_channels, stride=1, downsample=None):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = conv3x3(in_channels, out_channels, stride)\n",
    "        self.bn1 = nn.BatchNorm2d(out_channels)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = conv3x3(out_channels, out_channels)\n",
    "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
    "        self.downsample = downsample\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "        if self.downsample:\n",
    "            residual = self.downsample(x)\n",
    "        out += residual\n",
    "        out = self.relu(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# model\n",
    "class ResNet(nn.Module):\n",
    "    \"\"\"\n",
    "    Initialize  ResNet with forward propogation\n",
    "    \"\"\"\n",
    "    def __init__(self, block, layers, num_classes=10):\n",
    "        super(ResNet, self).__init__()\n",
    "        self.in_channels = 16\n",
    "        self.conv = conv3x3(3, 16)\n",
    "        self.bn = nn.BatchNorm2d(16)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.layer1 = self.make_layer(block, 16, layers[0])\n",
    "        self.layer2 = self.make_layer(block, 32, layers[1], 2)\n",
    "        self.layer3 = self.make_layer(block, 64, layers[2], 2)\n",
    "        self.avg_pool = nn.AvgPool2d(8)\n",
    "        self.fc = nn.Linear(64, num_classes)\n",
    "\n",
    "    def make_layer(self, block, out_channels, blocks, stride=1):\n",
    "        downsample = None\n",
    "        if (stride != 1) or (self.in_channels != out_channels):\n",
    "            downsample = nn.Sequential(\n",
    "                conv3x3(self.in_channels, out_channels, stride=stride),\n",
    "                nn.BatchNorm2d(out_channels))\n",
    "        layers = []\n",
    "        layers.append(block(self.in_channels, out_channels, stride, downsample))\n",
    "        self.in_channels = out_channels\n",
    "        for i in range(1, blocks):\n",
    "            layers.append(block(out_channels, out_channels))\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.conv(x)\n",
    "        out = self.bn(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.layer1(out)\n",
    "        out = self.layer2(out)\n",
    "        out = self.layer3(out)\n",
    "        out = self.avg_pool(out)\n",
    "        out = out.view(out.size(0), -1)\n",
    "        out = self.fc(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ResNet(ResidualBlock, [3, 3, 3]).to(device)\n",
    "\n",
    "# Define optimizer\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)\n",
    "\n",
    "# Define Cost function\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "# data and transformations\n",
    "\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])\n",
    "\n",
    "train_transform = transforms.Compose([\n",
    "    transforms.Pad(4),\n",
    "    transforms.RandomRotation(10),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomCrop(32),\n",
    "    transforms.ToTensor(),\n",
    "    normalize])\n",
    "    \n",
    "test_transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    normalize])\n",
    "\n",
    "\n",
    "train_dataset = torchvision.datasets.CIFAR10(root='data/',train=True, transform=train_transform,download=True)\n",
    "test_dataset = torchvision.datasets.CIFAR10(root='data/',train=False, transform=test_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# getting value of N\n",
    "N =len(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# important parameters\n",
    "\n",
    "sensitivity=1.0 # sensitivity is 1\n",
    "\n",
    "sigma_N=0.02*N # noise multiplier for N\n",
    "\n",
    "sigma_K=0.02*N # noise multiplier of K (sum of grad norms)\n",
    "\n",
    "C=0.1  # given innreference [2]\n",
    "\n",
    "C_star=4*C # External bound given as per paper\n",
    "\n",
    "gl=0.0001 # gl\n",
    "\n",
    "k_constant=5 # k in line 11, 24\n",
    "\n",
    "b=16 # expected batch size\n",
    "\n",
    "quantile=1 # lambda\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to add noise to a scalar\n",
    "\n",
    "def add_gaussian_noise(scalar, noise_multiplier):\n",
    "    \"\"\"\n",
    "    Adds Gaussian noise to a scalar.\n",
    "    \n",
    "    Parameters:\n",
    "    - scalar: The scalar value to which noise will be added.\n",
    "    - noise_multiplier: A factor to control the standard deviation of the Gaussian noise.\n",
    "    \n",
    "    Returns:\n",
    "    - The scalar value with added Gaussian noise.\n",
    "    \"\"\"\n",
    "\n",
    "    sigma=noise_multiplier*sensitivity\n",
    "\n",
    "    # Generate Gaussian noise with mean=0 and std=noise_multiplier\n",
    "    noise = torch.normal(mean=0.0, std=sigma,size=())\n",
    "    \n",
    "    # Add the noise to the scalar\n",
    "    noisy_scalar = scalar + noise\n",
    "    \n",
    "    return noisy_scalar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_tilda=add_gaussian_noise(N,sigma_N) # Noisy dataset size\n",
    "\n",
    "epochs=40 3 epoch number\n",
    "\n",
    "iterations=math.ceil(N/b) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[16], line 5\u001b[0m\n\u001b[0;32m      2\u001b[0m data_structure \u001b[38;5;241m=\u001b[39m []\n\u001b[0;32m      4\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mlen\u001b[39m(train_dataset)):\n\u001b[1;32m----> 5\u001b[0m    image_tensor, label \u001b[38;5;241m=\u001b[39m \u001b[43mtrain_dataset\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m\n\u001b[0;32m      6\u001b[0m    data_structure\u001b[38;5;241m.\u001b[39mappend([image_tensor, label, \u001b[38;5;241m0.0\u001b[39m, \u001b[38;5;241m0.0\u001b[39m,\u001b[38;5;241m0.0\u001b[39m,\u001b[38;5;241m0.0\u001b[39m,\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1.0\u001b[39m,i])\n\u001b[0;32m      8\u001b[0m K_cap_star\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m\n",
      "File \u001b[1;32mc:\\Users\\David\\anaconda3\\envs\\torch_env\\Lib\\site-packages\\torchvision\\datasets\\cifar.py:118\u001b[0m, in \u001b[0;36mCIFAR10.__getitem__\u001b[1;34m(self, index)\u001b[0m\n\u001b[0;32m    115\u001b[0m img \u001b[38;5;241m=\u001b[39m Image\u001b[38;5;241m.\u001b[39mfromarray(img)\n\u001b[0;32m    117\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransform \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m--> 118\u001b[0m     img \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtransform\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimg\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    120\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtarget_transform \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    121\u001b[0m     target \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtarget_transform(target)\n",
      "File \u001b[1;32mc:\\Users\\David\\anaconda3\\envs\\torch_env\\Lib\\site-packages\\torchvision\\transforms\\transforms.py:95\u001b[0m, in \u001b[0;36mCompose.__call__\u001b[1;34m(self, img)\u001b[0m\n\u001b[0;32m     93\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, img):\n\u001b[0;32m     94\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m t \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransforms:\n\u001b[1;32m---> 95\u001b[0m         img \u001b[38;5;241m=\u001b[39m \u001b[43mt\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimg\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     96\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m img\n",
      "File \u001b[1;32mc:\\Users\\David\\anaconda3\\envs\\torch_env\\Lib\\site-packages\\torchvision\\transforms\\transforms.py:137\u001b[0m, in \u001b[0;36mToTensor.__call__\u001b[1;34m(self, pic)\u001b[0m\n\u001b[0;32m    129\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, pic):\n\u001b[0;32m    130\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m    131\u001b[0m \u001b[38;5;124;03m    Args:\u001b[39;00m\n\u001b[0;32m    132\u001b[0m \u001b[38;5;124;03m        pic (PIL Image or numpy.ndarray): Image to be converted to tensor.\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    135\u001b[0m \u001b[38;5;124;03m        Tensor: Converted image.\u001b[39;00m\n\u001b[0;32m    136\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m--> 137\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto_tensor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpic\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\David\\anaconda3\\envs\\torch_env\\Lib\\site-packages\\torchvision\\transforms\\functional.py:171\u001b[0m, in \u001b[0;36mto_tensor\u001b[1;34m(pic)\u001b[0m\n\u001b[0;32m    169\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m pic\u001b[38;5;241m.\u001b[39mmode \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m1\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m    170\u001b[0m     img \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m255\u001b[39m \u001b[38;5;241m*\u001b[39m img\n\u001b[1;32m--> 171\u001b[0m img \u001b[38;5;241m=\u001b[39m \u001b[43mimg\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mview\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpic\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msize\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpic\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msize\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mF_pil\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_image_num_channels\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpic\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    172\u001b[0m \u001b[38;5;66;03m# put it from HWC to CHW format\u001b[39;00m\n\u001b[0;32m    173\u001b[0m img \u001b[38;5;241m=\u001b[39m img\u001b[38;5;241m.\u001b[39mpermute((\u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m1\u001b[39m))\u001b[38;5;241m.\u001b[39mcontiguous()\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for e in range(epochs):\n",
    "     data_structure = []\n",
    "\n",
    "     for i in range(len(train_dataset)):\n",
    "        image_tensor, label = train_dataset[i]\n",
    "        data_structure.append([image_tensor, label, 0.0, 0.0,0.0,0.0,-1.0,i]) # list to store the variablles for all data points\n",
    "\n",
    "     K_cap_star=0\n",
    "     model.train()\n",
    "     K=0\n",
    "     K_cap_list=[]\n",
    "     K_cap=0\n",
    "    \n",
    "     for t in range(math.ceil(iterations)):\n",
    "\n",
    "         if t==0:\n",
    "             \n",
    "             for i in range(len(data_structure)):\n",
    "                 \n",
    "                 image=data_structure[i][0]\n",
    "                 label = data_structure[i][1]\n",
    "                 image = image.unsqueeze(0)  \n",
    "                #  label = label.unsqueeze(0)\n",
    "                \n",
    "                 image=image.to(device)\n",
    "\n",
    "                 label = torch.tensor(label, dtype=torch.long)\n",
    "                 label = label.unsqueeze(0) \n",
    "                 label=label.to(device)\n",
    "\n",
    "                 # Zero gradients before each forward pass for this data point\n",
    "                 model.zero_grad()\n",
    "\n",
    "                 # Forward pass for this specific data point\n",
    "                 outputs = model(image)\n",
    "\n",
    "                 # Compute loss for this data point\n",
    "                 loss = criterion(outputs, label)\n",
    "\n",
    "                  # Backward pass\n",
    "                 loss.backward()\n",
    "\n",
    "                 \n",
    "\n",
    "                 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1/C)\n",
    "\n",
    "                 grad_norm=0.0\n",
    "\n",
    "                 for param in model.parameters():\n",
    "                     if param.grad is not None:\n",
    "                        grad_norm += param.grad.norm().item()\n",
    "\n",
    "                \n",
    "\n",
    "                 data_structure[i][2]=grad_norm\n",
    "                 \n",
    "                 data_structure[i][4]=k_constant*max(grad_norm,gl)\n",
    "                 \n",
    "                 K+=grad_norm\n",
    "\n",
    "                 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1/C_star)\n",
    "\n",
    "                 grad_norm_star=0.0\n",
    "\n",
    "                 for param in model.parameters():\n",
    "                     if param.grad is not None:\n",
    "                        grad_norm_star += param.grad.norm().item()\n",
    "                 \n",
    "                 data_structure[i][3]=grad_norm_star\n",
    "\n",
    "             K_cap_list.append(add_gaussian_noise(K,sigma_K))\n",
    "\n",
    "             for idx in range(iterations-1):\n",
    "                 K_cap_list.append(add_gaussian_noise(K_cap_list[-1],sigma_K))\n",
    "\n",
    "            \n",
    "         for i in range(len(data_structure)):\n",
    "\n",
    "             data_structure[i][5]=min(b*data_structure[i][4]/K_cap_list[t],1)\n",
    "\n",
    "         p_q=[row[4] for row in data_structure]\n",
    "\n",
    "         probabilities_q = np.array(p_q) / np.sum(p_q)\n",
    "\n",
    "         chi_q_idxs = np.random.choice(len(data_structure), size=k_constant*b, replace=False, p=probabilities_q)\n",
    "\n",
    "\n",
    "         for j in chi_q_idxs:\n",
    "\n",
    "            image=data_structure[j][0]\n",
    "            label = data_structure[j][1]\n",
    "            image = image.unsqueeze(0)  \n",
    "            # label = label.unsqueeze(0)\n",
    "\n",
    "            image=image.to(device)\n",
    "            label = torch.tensor(label, dtype=torch.long)\n",
    "            label = label.unsqueeze(0) \n",
    "            label=label.to(device)\n",
    "\n",
    "            # Zero gradients before each forward pass for this data point\n",
    "            model.zero_grad()\n",
    "            # Forward pass for this specific data point\n",
    "            outputs = model(image)\n",
    "            # Compute loss for this data point\n",
    "            loss = criterion(outputs, label)\n",
    "             # Backward pass\n",
    "            loss.backward()\n",
    "\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1/min(abs(data_structure[j][4]),C))\n",
    "\n",
    "            grad_norm=0.0\n",
    "\n",
    "            for param in model.parameters():\n",
    "                     if param.grad is not None:\n",
    "                        grad_norm += param.grad.norm().item()\n",
    "\n",
    "            data_structure[j][2]=grad_norm\n",
    "\n",
    "            data_structure[j][6]=grad_norm/data_structure[j][4]\n",
    "\n",
    "            data_structure[j][4]=k_constant*max(data_structure[j][2],gl)\n",
    "            \n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1/C_star)\n",
    "\n",
    "            grad_norm_star=0.0\n",
    "\n",
    "            for param in model.parameters():\n",
    "                     if param.grad is not None:\n",
    "                        grad_norm += param.grad.norm().item()\n",
    "\n",
    "            data_structure[j][3]=grad_norm_star\n",
    "\n",
    "         chi_q = [row for row in data_structure if row[6] != -1.0]\n",
    "\n",
    "         p_p = np.array([row[6] for row in chi_q])\n",
    "\n",
    "         probabilities_p = p_p/ np.sum(p_p)\n",
    "\n",
    "         chi_p=np.random.choice(chi_q, size=b, replace=False, p=probabilities_p)\n",
    "\n",
    "\n",
    "     # compute gradient for selected batch\n",
    "         # line 26\n",
    "     # optimize (line 27)\n",
    "         print(\"Executed till here\")\n",
    "\n",
    "     # summate K_cap_star\n",
    "     for i in range(len(data_structure)):\n",
    "         K_cap_star+=data_structure[i][3]\n",
    "     \n",
    "     K_cap_star=add_gaussian_noise(K_cap_star,0.02*N*C)\n",
    "     C=(K_cap_star/N_tilda)*quantile\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, test_loader, device, epoch):\n",
    "    \"\"\"\n",
    "    Evaluate the model on the test set\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "    accuracy = 100 * correct / total\n",
    "    print(f'Test Accuracy: {accuracy:.2f}%')\n",
    "    \n",
    "    # Save model checkpoint (only from rank 0)\n",
    "    torch.save(model.module.state_dict(), f'model_weight/resnet_epoch_{epoch}_acc_{accuracy:.2f}.ckpt')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_env",
   "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}