memory_footprint.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Memory footprint estimation\n",
    "\n",
    "Below are calculations on how much GPU memory is approximately required to run the [DAVEnet model](https://github.com/dharwath/DAVEnet-pytorch) (Harwath et al. 2018)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For importing the models folder\n",
    "import sys\n",
    "sys.path.append('/m/home/home4/44/virkkua1/unix/PlacesAudio_project/DAVEnet')\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "from torchsummary import summary\n",
    "\n",
    "import models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model consists of two CNNs, an audio and an image branch. The audio branch is a 5-layer CNN that generates audio embeddings from spoken captions. The image branch is a standard VGG16 where the final maxpool is replaced with 2D convolutional layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Davenet(\n",
       "  (batchnorm1): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv1): Conv2d(1, 128, kernel_size=(40, 1), stride=(1, 1))\n",
       "  (conv2): Conv2d(128, 256, kernel_size=(1, 11), stride=(1, 1), padding=(0, 5))\n",
       "  (conv3): Conv2d(256, 512, kernel_size=(1, 17), stride=(1, 1), padding=(0, 8))\n",
       "  (conv4): Conv2d(512, 512, kernel_size=(1, 17), stride=(1, 1), padding=(0, 8))\n",
       "  (conv5): Conv2d(512, 1024, kernel_size=(1, 17), stride=(1, 1), padding=(0, 8))\n",
       "  (pool): MaxPool2d(kernel_size=(1, 3), stride=(1, 2), padding=(0, 1), dilation=1, ceil_mode=False)\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Davenet(nn.Module):\n",
    "    def __init__(self, embedding_dim=1024):\n",
    "        super(Davenet, self).__init__()\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.batchnorm1 = nn.BatchNorm2d(1)\n",
    "        self.conv1 = nn.Conv2d(1, 128, kernel_size=(40,1), stride=(1,1), padding=(0,0))\n",
    "        self.conv2 = nn.Conv2d(128, 256, kernel_size=(1,11), stride=(1,1), padding=(0,5))\n",
    "        self.conv3 = nn.Conv2d(256, 512, kernel_size=(1,17), stride=(1,1), padding=(0,8))\n",
    "        self.conv4 = nn.Conv2d(512, 512, kernel_size=(1,17), stride=(1,1), padding=(0,8))\n",
    "        self.conv5 = nn.Conv2d(512, embedding_dim, kernel_size=(1,17), stride=(1,1), padding=(0,8))\n",
    "        self.pool = nn.MaxPool2d(kernel_size=(1,3), stride=(1,2),padding=(0,1))\n",
    "\n",
    "    def forward(self, x):\n",
    "        print(x.shape)\n",
    "        if x.dim() == 3:\n",
    "            x = x.unsqueeze(1)\n",
    "        print(x.shape)\n",
    "        x = self.batchnorm1(x)\n",
    "        print(x.shape)\n",
    "        x = F.relu(self.conv1(x))\n",
    "        x = F.relu(self.conv2(x))\n",
    "        x = self.pool(x)\n",
    "        x = F.relu(self.conv3(x))\n",
    "        x = self.pool(x)\n",
    "        x = F.relu(self.conv4(x))\n",
    "        x = self.pool(x)\n",
    "        x = F.relu(self.conv5(x))\n",
    "        x = self.pool(x)\n",
    "        print(x.shape)\n",
    "        x = x.squeeze(2)\n",
    "        print(x.shape)\n",
    "        return x\n",
    "\n",
    "#audio_model = models.Davenet()\n",
    "audio_model = Davenet()\n",
    "audio_model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 40, 1024])\n",
      "torch.Size([2, 1, 40, 1024])\n",
      "torch.Size([2, 1, 40, 1024])\n",
      "torch.Size([2, 1024, 1, 64])\n",
      "torch.Size([2, 1024, 64])\n",
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 128, 1, 1024]           5,248\n",
      "            Conv2d-2         [-1, 256, 1, 1024]         360,704\n",
      "         MaxPool2d-3          [-1, 256, 1, 512]               0\n",
      "            Conv2d-4          [-1, 512, 1, 512]       2,228,736\n",
      "         MaxPool2d-5          [-1, 512, 1, 256]               0\n",
      "            Conv2d-6          [-1, 512, 1, 256]       4,456,960\n",
      "         MaxPool2d-7          [-1, 512, 1, 128]               0\n",
      "            Conv2d-8         [-1, 1024, 1, 128]       8,913,920\n",
      "         MaxPool2d-9          [-1, 1024, 1, 64]               0\n",
      "================================================================\n",
      "Total params: 15,965,568\n",
      "Trainable params: 15,965,568\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.16\n",
      "Forward/backward pass size (MB): 10.00\n",
      "Params size (MB): 60.90\n",
      "Estimated Total Size (MB): 71.06\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "summary(audio_model, (40, 1024))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VGG16(\n",
       "  (image_model): Sequential(\n",
       "    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (1): ReLU(inplace=True)\n",
       "    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (3): ReLU(inplace=True)\n",
       "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (6): ReLU(inplace=True)\n",
       "    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (8): ReLU(inplace=True)\n",
       "    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (11): ReLU(inplace=True)\n",
       "    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (13): ReLU(inplace=True)\n",
       "    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (15): ReLU(inplace=True)\n",
       "    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (18): ReLU(inplace=True)\n",
       "    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (20): ReLU(inplace=True)\n",
       "    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (22): ReLU(inplace=True)\n",
       "    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (25): ReLU(inplace=True)\n",
       "    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (27): ReLU(inplace=True)\n",
       "    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (29): ReLU(inplace=True)\n",
       "    (30): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image_model = models.VGG16(pretrained=True)\n",
    "image_model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parameter counts are computed by the torchsummary module, which produces Keras-like network summaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 224, 224]           1,792\n",
      "              ReLU-2         [-1, 64, 224, 224]               0\n",
      "            Conv2d-3         [-1, 64, 224, 224]          36,928\n",
      "              ReLU-4         [-1, 64, 224, 224]               0\n",
      "         MaxPool2d-5         [-1, 64, 112, 112]               0\n",
      "            Conv2d-6        [-1, 128, 112, 112]          73,856\n",
      "              ReLU-7        [-1, 128, 112, 112]               0\n",
      "            Conv2d-8        [-1, 128, 112, 112]         147,584\n",
      "              ReLU-9        [-1, 128, 112, 112]               0\n",
      "        MaxPool2d-10          [-1, 128, 56, 56]               0\n",
      "           Conv2d-11          [-1, 256, 56, 56]         295,168\n",
      "             ReLU-12          [-1, 256, 56, 56]               0\n",
      "           Conv2d-13          [-1, 256, 56, 56]         590,080\n",
      "             ReLU-14          [-1, 256, 56, 56]               0\n",
      "           Conv2d-15          [-1, 256, 56, 56]         590,080\n",
      "             ReLU-16          [-1, 256, 56, 56]               0\n",
      "        MaxPool2d-17          [-1, 256, 28, 28]               0\n",
      "           Conv2d-18          [-1, 512, 28, 28]       1,180,160\n",
      "             ReLU-19          [-1, 512, 28, 28]               0\n",
      "           Conv2d-20          [-1, 512, 28, 28]       2,359,808\n",
      "             ReLU-21          [-1, 512, 28, 28]               0\n",
      "           Conv2d-22          [-1, 512, 28, 28]       2,359,808\n",
      "             ReLU-23          [-1, 512, 28, 28]               0\n",
      "        MaxPool2d-24          [-1, 512, 14, 14]               0\n",
      "           Conv2d-25          [-1, 512, 14, 14]       2,359,808\n",
      "             ReLU-26          [-1, 512, 14, 14]               0\n",
      "           Conv2d-27          [-1, 512, 14, 14]       2,359,808\n",
      "             ReLU-28          [-1, 512, 14, 14]               0\n",
      "           Conv2d-29          [-1, 512, 14, 14]       2,359,808\n",
      "             ReLU-30          [-1, 512, 14, 14]               0\n",
      "           Conv2d-31         [-1, 1024, 14, 14]       4,719,616\n",
      "================================================================\n",
      "Total params: 19,434,304\n",
      "Trainable params: 19,434,304\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 219.73\n",
      "Params size (MB): 74.14\n",
      "Estimated Total Size (MB): 294.44\n",
      "----------------------------------------------------------------\n",
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "       BatchNorm2d-1          [-1, 1, 40, 1024]               2\n",
      "            Conv2d-2         [-1, 128, 1, 1024]           5,248\n",
      "            Conv2d-3         [-1, 256, 1, 1024]         360,704\n",
      "         MaxPool2d-4          [-1, 256, 1, 512]               0\n",
      "            Conv2d-5          [-1, 512, 1, 512]       2,228,736\n",
      "         MaxPool2d-6          [-1, 512, 1, 256]               0\n",
      "            Conv2d-7          [-1, 512, 1, 256]       4,456,960\n",
      "         MaxPool2d-8          [-1, 512, 1, 128]               0\n",
      "            Conv2d-9         [-1, 1024, 1, 128]       8,913,920\n",
      "        MaxPool2d-10          [-1, 1024, 1, 64]               0\n",
      "================================================================\n",
      "Total params: 15,965,570\n",
      "Trainable params: 15,965,570\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.16\n",
      "Forward/backward pass size (MB): 10.31\n",
      "Params size (MB): 60.90\n",
      "Estimated Total Size (MB): 71.37\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# The second parameter is input size.\n",
    "summary(image_model, (3, 224, 224))\n",
    "summary(audio_model, (40, 1024))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To estimate the total memory footprint of a model, we need to add together the input size, forward *and* backward pass sizes (AKA feature + gradient maps or channels) and parameters' size (AKA weights and bias). We can see above that torchsummary can do this for us. But to understand where the numbers come from, let's do a sanity check below and compute the estimates for the audio network ourselves.\n",
    "\n",
    "The input size is computed as a product of the input dimensions. For forward pass size, dimensions of each layer are multiplied and the resulting products are summed together. Backward pass size is the same as forward pass, so forward pass can be multiplied by two. The weights/parameters for each layer are a product of the kernel size and input and output depths (i.e. the number of feature maps/channels). The potential bias vector size is added on top of the product. Parameters are naturally computed only for layers which have them, in this case convolutional and batch norm layers. In the end, all parameters are summed together. Finally, to convert each result to megabytes, they need to be multiplied by 4 bytes (values are presumably stored as float32, so each value uses 32 bits i.e. 4 bytes) and then divided by the square of 1024 (torchsummary uses binary notation). \n",
    "\n",
    "$$\\text{Input:} \\ \\ \\frac{40*1024*4\\text{B}}{1024^2} = 0.15625 \\text{MB} \\approx 0.16 \\text{MB}$$\n",
    "\n",
    "$$\\text{Forward+backward pass:} \\ \\ \\frac{40*1024+128*1024+256*1024+256*512+512^2+2*512*256+512*128+1024*128+1024*64}{1024^2}*4\\text{B}*2 = 10.3125 \\text{MB} \\approx 10.31 \\text{MB}$$\n",
    "\n",
    "$$\\text{Parameters:} \\ \\ \\frac{2+1*128*40+128+128*256*11+256+256*512*17+512+512^2*17+512+512*1024*17+1024}{1024^2}*4\\text{B} \\approx 60.90 \\text{MB}$$\n",
    "\n",
    "However, these estimates do not take batch size into account. When computing on a GPU, one batch is loaded to memory at a time. So when the batch size is for example 10, input size needs to be multiplied by 10. The same applies to forward and backward pass sizes, because activation and gradient maps are different for each input item in a batch. Parameters stay the same across different inputs, so the parameter size can be left untouched.\n",
    "\n",
    "In the DAVEnet paper, authors write that they train the model with batch size 128. Supplying the `batch_size` argument to the `summary()` gives us the estimate of\n",
    "\n",
    "$$28273.64 \\ \\text{MB} = \\frac{28273.64 \\ \\text{MB}}{1024} \\approx 27.61 \\ \\text{GB} $$\n",
    "\n",
    "for the image model and\n",
    "\n",
    "$$1400.90 \\ \\text{MB} = \\frac{1400.90 \\ \\text{MB}}{1024} \\approx 1.37 \\ \\text{GB}$$\n",
    "\n",
    "for the audio model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1        [128, 64, 224, 224]           1,792\n",
      "              ReLU-2        [128, 64, 224, 224]               0\n",
      "            Conv2d-3        [128, 64, 224, 224]          36,928\n",
      "              ReLU-4        [128, 64, 224, 224]               0\n",
      "         MaxPool2d-5        [128, 64, 112, 112]               0\n",
      "            Conv2d-6       [128, 128, 112, 112]          73,856\n",
      "              ReLU-7       [128, 128, 112, 112]               0\n",
      "            Conv2d-8       [128, 128, 112, 112]         147,584\n",
      "              ReLU-9       [128, 128, 112, 112]               0\n",
      "        MaxPool2d-10         [128, 128, 56, 56]               0\n",
      "           Conv2d-11         [128, 256, 56, 56]         295,168\n",
      "             ReLU-12         [128, 256, 56, 56]               0\n",
      "           Conv2d-13         [128, 256, 56, 56]         590,080\n",
      "             ReLU-14         [128, 256, 56, 56]               0\n",
      "           Conv2d-15         [128, 256, 56, 56]         590,080\n",
      "             ReLU-16         [128, 256, 56, 56]               0\n",
      "        MaxPool2d-17         [128, 256, 28, 28]               0\n",
      "           Conv2d-18         [128, 512, 28, 28]       1,180,160\n",
      "             ReLU-19         [128, 512, 28, 28]               0\n",
      "           Conv2d-20         [128, 512, 28, 28]       2,359,808\n",
      "             ReLU-21         [128, 512, 28, 28]               0\n",
      "           Conv2d-22         [128, 512, 28, 28]       2,359,808\n",
      "             ReLU-23         [128, 512, 28, 28]               0\n",
      "        MaxPool2d-24         [128, 512, 14, 14]               0\n",
      "           Conv2d-25         [128, 512, 14, 14]       2,359,808\n",
      "             ReLU-26         [128, 512, 14, 14]               0\n",
      "           Conv2d-27         [128, 512, 14, 14]       2,359,808\n",
      "             ReLU-28         [128, 512, 14, 14]               0\n",
      "           Conv2d-29         [128, 512, 14, 14]       2,359,808\n",
      "             ReLU-30         [128, 512, 14, 14]               0\n",
      "           Conv2d-31        [128, 1024, 14, 14]       4,719,616\n",
      "================================================================\n",
      "Total params: 19,434,304\n",
      "Trainable params: 19,434,304\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 73.50\n",
      "Forward/backward pass size (MB): 28126.00\n",
      "Params size (MB): 74.14\n",
      "Estimated Total Size (MB): 28273.64\n",
      "----------------------------------------------------------------\n",
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "       BatchNorm2d-1         [128, 1, 40, 1024]               2\n",
      "            Conv2d-2        [128, 128, 1, 1024]           5,248\n",
      "            Conv2d-3        [128, 256, 1, 1024]         360,704\n",
      "         MaxPool2d-4         [128, 256, 1, 512]               0\n",
      "            Conv2d-5         [128, 512, 1, 512]       2,228,736\n",
      "         MaxPool2d-6         [128, 512, 1, 256]               0\n",
      "            Conv2d-7         [128, 512, 1, 256]       4,456,960\n",
      "         MaxPool2d-8         [128, 512, 1, 128]               0\n",
      "            Conv2d-9        [128, 1024, 1, 128]       8,913,920\n",
      "        MaxPool2d-10         [128, 1024, 1, 64]               0\n",
      "================================================================\n",
      "Total params: 15,965,570\n",
      "Trainable params: 15,965,570\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 20.00\n",
      "Forward/backward pass size (MB): 1320.00\n",
      "Params size (MB): 60.90\n",
      "Estimated Total Size (MB): 1400.90\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "summary(image_model, (3, 224, 224), batch_size=128)\n",
    "summary(audio_model, (40, 1024), batch_size=128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cosine annealing test\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainables = [p for p in audio_model.parameters() if p.requires_grad]\n",
    "optimizer = torch.optim.SGD(trainables, 0.1, momentum=0.9, weight_decay=5e-7)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 100, eta_min=1e-6)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}