learn_AMP.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cead298d-d217-4004-9a3b-131b5a27a9c1",
   "metadata": {},
   "source": [
    "## taken from : \n",
    "https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "330866d8-6999-43be-9292-24fdbac3e22b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, time , gc\n",
    "start_time = None\n",
    "\n",
    "def start_timer():\n",
    "    global start_time\n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()\n",
    "    torch.cuda.reset_max_memory_allocated()\n",
    "    torch.cuda.synchronize()\n",
    "    start_time = time.time()\n",
    "\n",
    "def end_timer_and_print(local_msg):\n",
    "    torch.cuda.synchronize()\n",
    "    end_time = time.time()\n",
    "    print(\"\\n\" + local_msg)\n",
    "    print(\"Total execution time = {:.3f} sec\".format(end_time - start_time))\n",
    "    print(\"Max memory used by tensors = {} Mbytes\".format(torch.cuda.max_memory_allocated()*(1e-6)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e60046c0-b1ef-428b-8395-20e7a7a5427b",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 2. simple network, dataset, loss and so on (basic setup)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "8bdb531e-6e77-4640-9efb-777a346209c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_model(in_size, out_size, num_layers):\n",
    "    layers = []\n",
    "    for _ in range(num_layers - 1):\n",
    "        layers.append(torch.nn.Linear(in_size, in_size))\n",
    "        layers.append(torch.nn.ReLU())\n",
    "    layers.append(torch.nn.Linear(in_size, out_size))\n",
    "    return torch.nn.Sequential(*tuple(layers)).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "b979bca9-2b29-43a4-a207-786ca3c35c8d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "batch_size = 512 # Try, for example, 128, 256, 513.\n",
    "in_size = 4096\n",
    "out_size = 4096\n",
    "num_layers = 3\n",
    "num_batches = 50\n",
    "epochs = 3\n",
    "\n",
    "# Creates data in default precision.\n",
    "# The same data is used for both default and mixed precision trials below.\n",
    "# You don't need to manually change inputs' dtype when enabling mixed precision.\n",
    "data = [torch.randn(batch_size, in_size, device=\"cuda\") for _ in range(num_batches)]\n",
    "targets = [torch.randn(batch_size, out_size, device=\"cuda\") for _ in range(num_batches)]\n",
    "\n",
    "loss_fn = torch.nn.MSELoss().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "8985b159-4720-420c-a0e8-48b740a4daad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Forward Pass : \n",
      "Total execution time = 0.558 sec\n",
      "Max memory used by tensors = 1501.708288 Mbytes\n"
     ]
    }
   ],
   "source": [
    "start_timer()\n",
    "net = make_model(in_size, out_size, num_layers)\n",
    "for input, target in zip(data, targets):\n",
    "    output = net(input)\n",
    "    \n",
    "end_timer_and_print(\"Forward Pass : \")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f50443b7-7d01-43df-8f52-16b7dd6a8074",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 3. Default Precision (no AMP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "7369f05f-86ad-4315-a921-9dc636d1d039",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Default precision:\n",
      "Total execution time = 1.860 sec\n",
      "Max memory used by tensors = 1703.084032 Mbytes\n"
     ]
    }
   ],
   "source": [
    "net = make_model(in_size, out_size, num_layers)\n",
    "opt = torch.optim.SGD(net.parameters(), lr=0.001)\n",
    "\n",
    "start_timer()\n",
    "for epoch in range(epochs):\n",
    "    for input, target in zip(data, targets):\n",
    "        output = net(input)\n",
    "        loss = loss_fn(output, target)\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "        opt.zero_grad() # set_to_none=True here can modestly improve performance\n",
    "end_timer_and_print(\"Default precision:\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "384456aa-337f-4de2-9044-13bf8dc9d58f",
   "metadata": {},
   "source": [
    "### 4. Now doing AMP"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4bc8014-a14a-4afe-8c76-4cde90eea9f7",
   "metadata": {},
   "source": [
    "#### 4.1 Adding Autocast\n",
    "automatically casts parts of the model to fp16\n",
    "\n",
    "```\n",
    "with torch.autocast(device_type='cuda', dtype=torch.float16):\n",
    "    output = net(input)\n",
    "    loss = loss_fn(output, target)    \n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "0a5816d7-fd8e-404d-b90f-5eaaf4acb446",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Autocast :\n",
      "Total execution time = 0.691 sec\n",
      "Max memory used by tensors = 1405.198848 Mbytes\n"
     ]
    }
   ],
   "source": [
    "start_timer()\n",
    "for epoch in range(epochs): # 0 epochs, this section is for illustration only\n",
    "    for input, target in zip(data, targets):\n",
    "        # Runs the forward pass under autocast.\n",
    "        with torch.autocast(device_type='cuda', dtype=torch.float16):\n",
    "            output = net(input)\n",
    "            # output is float16 because linear layers autocast to float16.\n",
    "            assert output.dtype is torch.float16\n",
    "\n",
    "            loss = loss_fn(output, target)\n",
    "            # loss is float32 because mse_loss layers autocast to float32.\n",
    "            assert loss.dtype is torch.float32\n",
    "\n",
    "        # Exits autocast before backward().\n",
    "        # Backward passes under autocast are not recommended.\n",
    "        # Backward ops run in the same dtype autocast chose for corresponding forward ops.\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "        opt.zero_grad() # set_to_none=True here can modestly improve performance\n",
    "end_timer_and_print(\"Autocast :\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc20f6fa-564e-4ad6-8b31-1fa4ce714d42",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### 4.2. Adding GradScaler\n",
    "\n",
    "scaling the grad by S and so on (prevent small grad values from overflow)\n",
    "\n",
    "=> MUST CREATE NEW GRAD FOR EVERY NEW FOLD!  \n",
    "\n",
    "```\n",
    "scaler = torch.cuda.amp.GradScaler() #an object\n",
    "##do the autocast we did before \n",
    "with torch.autocast(device_type='cuda', dtype=torch.float16):\n",
    "    output = net(input)\n",
    "    loss = loss_fn(output, target)\n",
    "\n",
    "scaler.scale(loss).backward() #do backward ON the scaled (fp16) loss \n",
    "\n",
    "scaler.step(opt) #scale the opt (unscales the gradients and if it does not contian NaN do actual optimizer.step() \n",
    "\n",
    "scaler.update() #update the scaler => updates loss and opt simulatenously I think\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "b5abff8a-93db-48b5-bcd1-a26a0a1e7e9d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Autocast and GradScaler :\n",
      "Total execution time = 0.795 sec\n",
      "Max memory used by tensors = 1405.200384 Mbytes\n"
     ]
    }
   ],
   "source": [
    "start_timer()\n",
    "# Constructs scaler once, at the beginning of the convergence run, using default args.\n",
    "# If your network fails to converge with default GradScaler args, please file an issue.\n",
    "# The same GradScaler instance should be used for the entire convergence run.\n",
    "# If you perform multiple convergence runs in the same script, each run should use\n",
    "# a dedicated fresh GradScaler instance.  GradScaler instances are lightweight.\n",
    "scaler = torch.cuda.amp.GradScaler() \n",
    "\n",
    "for epoch in range(epochs): # 0 epochs, this section is for illustration only\n",
    "    for input, target in zip(data, targets):\n",
    "        with torch.autocast(device_type='cuda', dtype=torch.float16):\n",
    "            output = net(input)\n",
    "            loss = loss_fn(output, target)\n",
    "\n",
    "        # Scales loss.  Calls backward() on scaled loss to create scaled gradients.\n",
    "        scaler.scale(loss).backward()\n",
    "\n",
    "        # scaler.step() first unscales the gradients of the optimizer's assigned params.\n",
    "        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,\n",
    "        # otherwise, optimizer.step() is skipped.\n",
    "        scaler.step(opt)\n",
    "\n",
    "        # Updates the scale for next iteration.\n",
    "        scaler.update()\n",
    "\n",
    "        opt.zero_grad() # set_to_none=True here can modestly improve performance\n",
    "end_timer_and_print(\"Autocast and GradScaler :\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3880413-f80a-4d44-bdd2-9fd3419bd3d7",
   "metadata": {},
   "source": [
    "#### 4.3 easily switching b/w AMP and default \n",
    "\n",
    "`GradScaler` 과 `autocast`에서 (`enabled = True or False`)\n",
    "\n",
    " ```\n",
    " scaler = torch.cuda.amp.GradScaler(enabled = enable_AMP) \n",
    " with torch.autocast(device_type='cuda', dtype=torch.float16, enabled = enable_AMP):\n",
    " ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "c147aabd-a64b-4c8c-96eb-1ec8fa037ebc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Autocast and GradScaler :\n",
      "Total execution time = 1.859 sec\n",
      "Max memory used by tensors = 1359.069696 Mbytes\n"
     ]
    }
   ],
   "source": [
    "start_timer()\n",
    "enable_AMP = False\n",
    "scaler = torch.cuda.amp.GradScaler(enabled = enable_AMP) \n",
    "\n",
    "for epoch in range(epochs): # 0 epochs, this section is for illustration only\n",
    "    for input, target in zip(data, targets):\n",
    "        with torch.autocast(device_type='cuda', dtype=torch.float16, enabled = enable_AMP):\n",
    "            output = net(input)\n",
    "            loss = loss_fn(output, target)\n",
    "\n",
    "        # Scales loss.  Calls backward() on scaled loss to create scaled gradients.\n",
    "        scaler.scale(loss).backward()\n",
    "\n",
    "        # scaler.step() first unscales the gradients of the optimizer's assigned params.\n",
    "        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,\n",
    "        # otherwise, optimizer.step() is skipped.\n",
    "        scaler.step(opt)\n",
    "\n",
    "        # Updates the scale for next iteration.\n",
    "        scaler.update()\n",
    "\n",
    "        opt.zero_grad() # set_to_none=True here can modestly improve performance\n",
    "end_timer_and_print(\"Autocast and GradScaler :\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c48479a8-c4de-4322-acab-1a4bcaae22d7",
   "metadata": {},
   "source": [
    "나머지 : 생략 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be2e8ab7-327a-4a3a-bd54-15dae3543526",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "VAE_3DCNN_older_MONAI",
   "language": "python",
   "name": "vae_3dcnn_older_monai"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}