12_Memory_management_and_Multi_GPU_Usage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "781ccfd5",
   "metadata": {},
   "source": [
    "In this notebook, we'll cover the following topics divided into sections:\n",
    "\n",
    "**1: Utilizing Multiple GPUs**\n",
    "We'll explore how to leverage multiple GPUs for your neural network, employing either data parallelism or model parallelism for efficient training.\n",
    "\n",
    "**2: Automating GPU Selection**\n",
    "We'll discuss techniques to automate GPU selection when creating new objects, ensuring efficient resource utilization in systems with multiple GPUs.\n",
    "\n",
    "**3: Diagnosing Memory Issues**\n",
    "We'll cover techniques for diagnosing and analyzing memory-related issues that may arise during model training, providing insights into effective memory utilization and troubleshooting common problems.\n",
    "\n",
    "#### Moving Tensors between CPUs and GPUs\n",
    "To facilitate tensor movement between CPUs and GPUs in PyTorch, the to() or cuda() function is utilized. By providing the index of the GPU as an argument, tensors can be efficiently moved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4d8e5f57",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "05bd6c13",
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    dev = \"cuda:0\"\n",
    "else:\n",
    "    dev = \"cpu\"\n",
    "\n",
    "device = torch.device(dev)\n",
    "\n",
    "a = torch.zeros(4, 3)\n",
    "a = a.to(0)  # alternatively, a.to(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f57fc595",
   "metadata": {},
   "source": [
    "Similarly, this functionality extends to moving `nn.Module` objects between GPUs as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6a0009f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "myNetwork(\n",
       "  (net): Linear(in_features=5, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class myNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.net = nn.Linear(5, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "clf = myNetwork()\n",
    "clf.to(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b99e0164",
   "metadata": {},
   "source": [
    "#### Retrieving Device Information\n",
    "To retrieve the device of a tensor, the get_device() method can be employed. It's important to note that this method is only supported for GPU tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c6c32c69",
   "metadata": {},
   "outputs": [],
   "source": [
    "dev = a.get_device()\n",
    "b = torch.tensor(a.shape).to(dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "590a21c0",
   "metadata": {},
   "source": [
    "Additionally, the default device for creating GPU tensors can be set.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9aea57b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.set_device(0)\n",
    "\n",
    "tens = torch.Tensor(3, 4).cuda()\n",
    "tens.get_device()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a83d4fa3",
   "metadata": {},
   "source": [
    "#### The new_* functions\n",
    "\n",
    "The `new_` functions in PyTorch, introduced in version `1.0`, offer convenient ways to create new tensors based on existing tensors. When a function like `new_ones` is called on a Tensor, it returns a new tensor of the same data type and on the same device as the tensor on which the `new_ones` function was invoked."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "18b36001",
   "metadata": {},
   "outputs": [],
   "source": [
    "ones = torch.ones((2,)).cuda(0)\n",
    "\n",
    "# Create a tensor of ones of size (3,4) on the same device as \"ones\"\n",
    "newOnes = ones.new_ones((3,4)) \n",
    "\n",
    "randTensor = torch.randn(2,4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e32a24d",
   "metadata": {},
   "source": [
    "These `new_` functions provide a quick and efficient way to generate new tensors with desired properties while maintaining consistency with existing tensors. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f0afb58",
   "metadata": {},
   "source": [
    "#### Using Multiple GPUs\n",
    "\n",
    "There are two primary ways to utilize multiple GPUs in PyTorch:\n",
    "\n",
    "**Data Parallelism**\n",
    "Data Parallelism involves dividing batches into smaller batches and processing them in parallel across multiple GPUs. In PyTorch, this is achieved through the `nn.DataParallel` class. You initialize a `nn.DataParallel` object with an `nn.Module` object representing your network and a list of GPU IDs across which the batches are to be parallelized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "69c3aeec",
   "metadata": {},
   "outputs": [],
   "source": [
    "myNet = myNetwork()\n",
    "parallel_net = nn.DataParallel(myNet, device_ids=[0])\n",
    "\n",
    "inputs = torch.Tensor(5,)    # random inputs \n",
    "inputs = inputs.to(0)\n",
    "myNet.to(0)\n",
    "\n",
    "predictions = parallel_net(inputs)\n",
    "loss = (1 - predictions).mean()\n",
    "loss.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d26a1280",
   "metadata": {},
   "source": [
    "One issue with DataParallel is that it can lead to asymmetrical load on one GPU (the main node). This can be mitigated by computing the loss during the forward pass or by implementing a parallel loss function layer.\n",
    "\n",
    "#### Model Parallelism\n",
    "Model Parallelism involves breaking the neural network into smaller sub-networks and executing these sub-networks on different GPUs. This approach is useful when the network is too large to fit inside a single GPU. Implementing Model Parallelism in PyTorch is straightforward:\n",
    "\n",
    "1. Ensure the input and the network are always on the same device.\n",
    "\n",
    "2. Utilize the `to()` and `cuda()` functions with autograd support for gradient copying between GPUs during the backward pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "88d9b5a4",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA error: invalid device ordinal\nCUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1.",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[10], line 20\u001b[0m\n\u001b[0;32m     17\u001b[0m x \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mTensor(\u001b[38;5;241m100\u001b[39m,)        \u001b[38;5;66;03m# Random Input\u001b[39;00m\n\u001b[0;32m     18\u001b[0m x \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39mto(\u001b[38;5;241m0\u001b[39m)\n\u001b[1;32m---> 20\u001b[0m net \u001b[38;5;241m=\u001b[39m \u001b[43mmodel_parallel\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m        \u001b[38;5;66;03m# No need to put it on GPUs as that has been taken care of in the init function\u001b[39;00m\n\u001b[0;32m     22\u001b[0m loss \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m1\u001b[39m \u001b[38;5;241m-\u001b[39m net(x))\u001b[38;5;241m.\u001b[39mmean()\n\u001b[0;32m     23\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n",
      "Cell \u001b[1;32mIn[10], line 8\u001b[0m, in \u001b[0;36mmodel_parallel.__init__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msub_network2 \u001b[38;5;241m=\u001b[39m nn\u001b[38;5;241m.\u001b[39mLinear(\u001b[38;5;241m32\u001b[39m, \u001b[38;5;241m10\u001b[39m)     \u001b[38;5;66;03m# This part stays on GPU 2\u001b[39;00m\n\u001b[0;32m      7\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msub_network1\u001b[38;5;241m.\u001b[39mcuda(\u001b[38;5;241m0\u001b[39m)\n\u001b[1;32m----> 8\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msub_network2\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcuda\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\.conda\\envs\\pytorch\\lib\\site-packages\\torch\\nn\\modules\\module.py:749\u001b[0m, in \u001b[0;36mModule.cuda\u001b[1;34m(self, device)\u001b[0m\n\u001b[0;32m    732\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcuda\u001b[39m(\u001b[38;5;28mself\u001b[39m: T, device: Optional[Union[\u001b[38;5;28mint\u001b[39m, device]] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[0;32m    733\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Moves all model parameters and buffers to the GPU.\u001b[39;00m\n\u001b[0;32m    734\u001b[0m \n\u001b[0;32m    735\u001b[0m \u001b[38;5;124;03m    This also makes associated parameters and buffers different objects. So\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    747\u001b[0m \u001b[38;5;124;03m        Module: self\u001b[39;00m\n\u001b[0;32m    748\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m--> 749\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_apply\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcuda\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\.conda\\envs\\pytorch\\lib\\site-packages\\torch\\nn\\modules\\module.py:664\u001b[0m, in \u001b[0;36mModule._apply\u001b[1;34m(self, fn)\u001b[0m\n\u001b[0;32m    660\u001b[0m \u001b[38;5;66;03m# Tensors stored in modules are graph leaves, and we don't want to\u001b[39;00m\n\u001b[0;32m    661\u001b[0m \u001b[38;5;66;03m# track autograd history of `param_applied`, so we have to use\u001b[39;00m\n\u001b[0;32m    662\u001b[0m \u001b[38;5;66;03m# `with torch.no_grad():`\u001b[39;00m\n\u001b[0;32m    663\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[1;32m--> 664\u001b[0m     param_applied \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mparam\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    665\u001b[0m should_use_set_data \u001b[38;5;241m=\u001b[39m compute_should_use_set_data(param, param_applied)\n\u001b[0;32m    666\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m should_use_set_data:\n",
      "File \u001b[1;32m~\\.conda\\envs\\pytorch\\lib\\site-packages\\torch\\nn\\modules\\module.py:749\u001b[0m, in \u001b[0;36mModule.cuda.<locals>.<lambda>\u001b[1;34m(t)\u001b[0m\n\u001b[0;32m    732\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcuda\u001b[39m(\u001b[38;5;28mself\u001b[39m: T, device: Optional[Union[\u001b[38;5;28mint\u001b[39m, device]] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[0;32m    733\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Moves all model parameters and buffers to the GPU.\u001b[39;00m\n\u001b[0;32m    734\u001b[0m \n\u001b[0;32m    735\u001b[0m \u001b[38;5;124;03m    This also makes associated parameters and buffers different objects. So\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    747\u001b[0m \u001b[38;5;124;03m        Module: self\u001b[39;00m\n\u001b[0;32m    748\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m--> 749\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_apply(\u001b[38;5;28;01mlambda\u001b[39;00m t: \u001b[43mt\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcuda\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m)\n",
      "\u001b[1;31mRuntimeError\u001b[0m: CUDA error: invalid device ordinal\nCUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1."
     ]
    }
   ],
   "source": [
    "class model_parallel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.sub_network1 = nn.Linear(100, 32)    # This part stays on GPU 1\n",
    "        self.sub_network2 = nn.Linear(32, 10)     # This part stays on GPU 2\n",
    "\n",
    "        self.sub_network1.cuda(0)\n",
    "        self.sub_network2.cuda(1)\n",
    "\n",
    "    def forward(x):\n",
    "        x = x.cuda(0)\n",
    "        x = self.sub_network1(x)\n",
    "        x = x.cuda(1)\n",
    "        x = self.sub_network2(x)\n",
    "        return x\n",
    "\n",
    "x = torch.Tensor(100,)        # Random Input\n",
    "x = x.to(0)\n",
    "\n",
    "net = model_parallel()        # No need to put it on GPUs as that has been taken care of in the init function\n",
    "\n",
    "loss = (1 - net(x)).mean()\n",
    "loss.backward()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28b32572",
   "metadata": {},
   "source": [
    "These methods enable efficient utilization of multiple GPUs in PyTorch, offering scalability and performance improvements for deep learning tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33d4b6bc",
   "metadata": {},
   "source": [
    "#### Dealing with Memory Losses using `del` keyword\n",
    "\n",
    "While PyTorch has an aggressive garbage collector, it frees up the variable only when there exist no Pythonic reference to the object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "783904da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.0400,  0.6801,  0.3101,  1.3898]])\n"
     ]
    }
   ],
   "source": [
    "for x in range(10):\n",
    "    tensor = torch.randn(1, 4)\n",
    "\n",
    "print(tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f643f784",
   "metadata": {},
   "source": [
    "A good practice to get rid of these variables is by using the del keyword.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2c07aba8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([0.0464], grad_fn=<AddBackward0>) tensor(0.9536, grad_fn=<MeanBackward0>)\n"
     ]
    }
   ],
   "source": [
    "net = myNetwork()\n",
    "opt = torch.optim.SGD(net.parameters(), lr=0.01)\n",
    "\n",
    "inp = torch.randn(5,)\n",
    "\n",
    "for x in range(10):\n",
    "    out = net(inp)\n",
    "    loss = (1 - out).mean()\n",
    "    opt.zero_grad()\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "\n",
    "print(out, loss)                      # these variables still exist\n",
    "del out, loss                         # Free the memory taken by these variables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6adc7889",
   "metadata": {},
   "source": [
    "#### Using Python Data Types Instead Of 1-D Tensors\n",
    "\n",
    "Often, we aggregate values in our training loop to compute some metrics. If not done carefully in PyTorch, such a thing can lead to excess use of memory than what is required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a8eb51cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_loss = 0\n",
    "\n",
    "for x in range(10):\n",
    "    iter_loss = torch.randn(3, 4).mean()\n",
    "    iter_loss.requires_grad = True     # losses are supposed to be differentiable\n",
    "    total_loss += iter_loss            # use total_loss += iter_loss.item() instead"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9f634c3",
   "metadata": {},
   "source": [
    "#### Using `torch.no_grad()` for inference\n",
    "Whenever you are doing inference with your network or any operation that doesn't require backpropagation of gradients, you should always put the code inside torch.no_grad() context manager."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ee17e32f",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = myNetwork()\n",
    "inp = torch.randn(5,)\n",
    "\n",
    "with torch.no_grad():\n",
    "    out = net(inp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39093148",
   "metadata": {},
   "source": [
    "#### Emptying CUDA cache\n",
    "While PyTorch aggressively frees up memory, a PyTorch process may not give back the memory to the OS even after you delete your tensors. This memory is cached so that it can be quickly allocated to new tensors being allocated without requesting extra memory from the OS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "bb03dabb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial GPU Usage\n",
      "| ID | GPU | MEM |\n",
      "------------------\n",
      "|  0 |  6% |  6% |\n",
      "GPU Usage after allocating a bunch of Tensors\n",
      "| ID | GPU | MEM |\n",
      "------------------\n",
      "|  0 |  0% | 53% |\n",
      "GPU Usage after deleting the Tensors\n",
      "| ID | GPU | MEM |\n",
      "------------------\n",
      "|  0 | 28% | 53% |\n",
      "GPU Usage after emptying the cache\n",
      "| ID | GPU | MEM |\n",
      "------------------\n",
      "|  0 | 28% |  6% |\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from GPUtil import showUtilization as gpu_usage\n",
    "\n",
    "print(\"Initial GPU Usage\")\n",
    "gpu_usage()                             \n",
    "\n",
    "tensorList = []\n",
    "for x in range(10):\n",
    "    tensorList.append(torch.randn(10000000, 10).cuda())   # reduce the size of tensor if you are getting OOM\n",
    "\n",
    "print(\"GPU Usage after allocating a bunch of Tensors\")\n",
    "gpu_usage()\n",
    "\n",
    "del tensorList\n",
    "\n",
    "print(\"GPU Usage after deleting the Tensors\")\n",
    "gpu_usage()  \n",
    "\n",
    "print(\"GPU Usage after emptying the cache\")\n",
    "torch.cuda.empty_cache()\n",
    "gpu_usage()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "618365fc",
   "metadata": {},
   "source": [
    "#### Using CUDNN Backend\n",
    "One can use the CUDNN benchmark to have optimizations in the code. These optimizations are especially beneficial if your input size is fixed (you are not using RNNs)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f62c7195",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "torch.backends.cudnn.benchmark = True\n",
    "torch.backends.cudnn.enabled = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53ca53c4",
   "metadata": {},
   "source": [
    "#### Using Half Precision Floats\n",
    "\n",
    "One can use half precision floats if the GPU has FP16 support. It's simple enough to convert a normal model to its half precision variant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "48cbfd58",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.6973], device='cuda:0', dtype=torch.float16, grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inp = torch.randn(5,).cuda().half()\n",
    "\n",
    "model = myNetwork().cuda().half()\n",
    "\n",
    "model(inp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "783382ca",
   "metadata": {},
   "source": [
    "Batch Norm layers have been reported to have convergence issues with half precision floats, so it's better to use full precision for them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "c1c4e0d4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.2389]], device='cuda:0', dtype=torch.float16,\n",
       "       grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class myNetworkBN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.l1 = nn.Linear(10,5)\n",
    "        self.bn = nn.BatchNorm1d(5)\n",
    "        self.l2 = nn.Linear(5,1)\n",
    "     \n",
    "    def forward(self,x):\n",
    "        x = self.l1(x)\n",
    "        x = self.bn(x)\n",
    "        x = self.l2(x)\n",
    "        return x \n",
    "\n",
    "inp = torch.randn(10,).cuda().half().unsqueeze(0)       # Unsqueeze op to add mini-batch dimension\n",
    "\n",
    "model = myNetworkBN().cuda().half().eval()              # Eval mode = use population statistics in BN\n",
    "\n",
    "model(inp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c2d7acf",
   "metadata": {},
   "source": [
    "These techniques help manage memory efficiently in PyTorch, ensuring optimal utilization of system resources. One must always be careful about half precision floats when the value may get too large. It is recommended to use the Nvidia `apex` extension for using mixed precision training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "250dfa1a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}