mem_leak.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as T\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import numpy as np\n",
    "from __future__ import print_function\n",
    "import gc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$x_{t+1} = (A + pB)x_t + Cu_{t+1}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyData(Dataset):    \n",
    "    def __init__(self, u, p, x,n_future_steps=1):\n",
    "        self.x = nn.Parameter(x,requires_grad=False)\n",
    "        self.p = nn.Parameter(p,requires_grad=False)\n",
    "        self.u = nn.Parameter(u,requires_grad=False)\n",
    "        self.nfeatures = x.shape[1]\n",
    "        self.n_future_steps = n_future_steps\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.x)-self.n_future_steps\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        indices = slice(idx,idx+self.n_future_steps)\n",
    "        x_true_indices = slice(idx+1,idx+self.n_future_steps+1)\n",
    "        return (self.u[indices], self.p[indices],\n",
    "                self.x[indices], self.x[x_true_indices])\n",
    "\n",
    "class Dynamics(nn.Module):\n",
    "    def __init__(self, nfeatures):\n",
    "        super(Dynamics, self).__init__()\n",
    "        self.A = nn.Parameter(T.normal(T.zeros(nfeatures,nfeatures),0.5),requires_grad=True)\n",
    "        self.B = nn.Parameter(T.normal(T.zeros(nfeatures,nfeatures),0.5),requires_grad=True)\n",
    "        self.C = nn.Parameter(T.normal(T.zeros(nfeatures),0.5),requires_grad=True)\n",
    "\n",
    "    def forward(self, u, p, x):\n",
    "        return (x[:,0] + (T.matmul((self.A + p[:,0,None,None]*self.B), x[:,0,:,None]).squeeze()) + u[:,0,None] * self.C)[:,None]\n",
    "\n",
    "\n",
    "def train(model,data,nepochs=10, lambdaA=1e-8, lambdaB=1e-6, lr=0.001):\n",
    "    dataloader = DataLoader(data, batch_size=batch_size, shuffle=True)\n",
    "    optimizer = T.optim.Adam(model.parameters(),lr=lr)\n",
    "    og_mem = T.cuda.memory_allocated() / 1024**2\n",
    "    print(\"Allocated Memory: {} MB\".format(og_mem))\n",
    "    for e in range(nepochs):\n",
    "        for batch_data in dataloader:\n",
    "            U,P,X, X_true = batch_data\n",
    "            X_pred = model(U,P,X)\n",
    "            mse_loss = F.mse_loss(X_pred,X_true)\n",
    "            l1_B = model.B.norm(1)\n",
    "            l1_A = model.A.norm(1)\n",
    "            \n",
    "            loss = mse_loss + lambdaA*l1_A + lambdaB*l1_B\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            del X_pred, U,P,X, X_true, mse_loss, l1_A, l1_B, loss\n",
    "            gc.collect()\n",
    "            torch.cuda.empty_cache()\n",
    "            \n",
    "            mem = T.cuda.memory_allocated() / 1024**2\n",
    "            print(\"New allocations: {} MB\".format(mem-og_mem))\n",
    "            og_mem = mem\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "ntime = 2826\n",
    "nstim = 30\n",
    "nfeatures = 15888\n",
    "\n",
    "u_train = T.rand(ntime).cuda()\n",
    "p_train = T.rand(ntime).cuda()\n",
    "time_train = T.from_numpy(np.arange(ntime)).cuda()\n",
    "x_train = T.rand(ntime,nfeatures).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dynamics()"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_future_steps = 1\n",
    "batch_size = 1\n",
    "\n",
    "data = MyData(u_train,p_train,x_train,n_future_steps)\n",
    "model = Dynamics(data.nfeatures)\n",
    "model.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Allocated Memory: 2098 MB\n",
      "New allocations: 0 MB\n",
      "New allocations: 0 MB\n",
      "New allocations: 0 MB\n",
      "New allocations: 0 MB\n",
      "New allocations: 0 MB\n",
      "New allocations: 0 MB\n",
      "New allocations: 0 MB\n",
      "New allocations: 0 MB\n",
      "New allocations: 0 MB\n",
      "New allocations: 0 MB\n"
     ]
    }
   ],
   "source": [
    "og_mem = T.cuda.memory_allocated() / 1024**2\n",
    "print(\"Allocated Memory: {} MB\".format(og_mem))\n",
    "with T.no_grad():\n",
    "    for t in range(10):\n",
    "        x_pred = T.squeeze(model(u_train[[t],None], p_train[[t],None], \n",
    "                             x_train[[t],None]))\n",
    "        mem = T.cuda.memory_allocated() / 1024**2\n",
    "        print(\"New allocations: {} MB\".format(mem-og_mem))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Allocated Memory: 2098 MB\n",
      "New allocations: 5778 MB\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1524577523076/work/aten/src/THC/generic/THCStorage.cu:58",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-6-f76860546b30>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1e-4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1e-6\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.001\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-2-d5d0d1975ac8>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(model, data, nepochs, lambdaA, lambdaB, lr)\u001b[0m\n\u001b[1;32m     42\u001b[0m             \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmse_loss\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mlambdaA\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0ml1_A\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mlambdaB\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0ml1_B\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     43\u001b[0m             \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 44\u001b[0;31m             \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     45\u001b[0m             \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     46\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/home/ubuntu/anaconda3/envs/pytorch_p27/lib/python2.7/site-packages/torch/tensor.pyc\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph)\u001b[0m\n\u001b[1;32m     91\u001b[0m                 \u001b[0mproducts\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0mDefaults\u001b[0m \u001b[0mto\u001b[0m \u001b[0;34m`\u001b[0m\u001b[0;34m`\u001b[0m\u001b[0mFalse\u001b[0m\u001b[0;34m`\u001b[0m\u001b[0;34m`\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     92\u001b[0m         \"\"\"\n\u001b[0;32m---> 93\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     94\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     95\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/home/ubuntu/anaconda3/envs/pytorch_p27/lib/python2.7/site-packages/torch/autograd/__init__.pyc\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables)\u001b[0m\n\u001b[1;32m     87\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m     88\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 89\u001b[0;31m         allow_unreachable=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m     90\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1524577523076/work/aten/src/THC/generic/THCStorage.cu:58"
     ]
    }
   ],
   "source": [
    "train(model,data,1,1e-4,1e-6,lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.Tensor'> torch.Size([1, 1])\n",
      "<class 'torch.Tensor'> torch.Size([1, 1])\n",
      "<class 'torch.Tensor'> torch.Size([1, 1, 15888])\n",
      "<class 'torch.Tensor'> torch.Size([1, 1, 15888])\n",
      "<class 'torch.Tensor'> torch.Size([1, 1, 15888])\n",
      "<class 'torch.Tensor'> torch.Size([])\n",
      "<class 'torch.Tensor'> torch.Size([])\n",
      "<class 'torch.Tensor'> torch.Size([])\n",
      "<class 'torch.Tensor'> torch.Size([])\n",
      "<class 'torch.Tensor'> torch.Size([])\n",
      "<class 'torch.Tensor'> torch.Size([2826])\n",
      "<class 'torch.Tensor'> torch.Size([2826])\n",
      "<class 'torch.Tensor'> torch.Size([2826])\n",
      "<class 'torch.Tensor'> torch.Size([2826, 15888])\n",
      "<class 'torch.nn.parameter.Parameter'> torch.Size([2826])\n",
      "<class 'torch.nn.parameter.Parameter'> torch.Size([2826])\n",
      "<class 'torch.nn.parameter.Parameter'> torch.Size([15888, 15888])\n",
      "<class 'torch.nn.parameter.Parameter'> torch.Size([15888, 15888])\n",
      "<class 'torch.nn.parameter.Parameter'> torch.Size([15888])\n",
      "<class 'torch.Tensor'> torch.Size([15888])\n",
      "<class 'torch.Tensor'> torch.Size([15888, 15888])\n",
      "<class 'torch.Tensor'> torch.Size([15888, 15888])\n",
      "<class 'torch.Tensor'> torch.Size([15888])\n",
      "<class 'torch.nn.parameter.Parameter'> torch.Size([2826, 15888])\n",
      "<class 'torch.Tensor'> torch.Size([15888, 15888])\n",
      "<class 'torch.Tensor'> torch.Size([15888, 15888])\n",
      "<class 'torch.Tensor'> torch.Size([15888])\n"
     ]
    }
   ],
   "source": [
    "for obj in gc.get_objects():\n",
    "    if T.is_tensor(obj) or (hasattr(obj, 'data') and T.is_tensor(obj.data)):\n",
    "        print(type(obj), obj.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_pytorch_p27)",
   "language": "python",
   "name": "conda_pytorch_p27"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}