optimizer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import torch \n",
    "from torch import nn, optim \n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import clear_output\n",
    "\n",
    "import time \n",
    "\n",
    "from lib import *\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class my_TT_linear(nn.Module):\n",
    "    def __init__(self, inp_modes, out_modes, ranks):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.d = len(inp_modes)\n",
    "        self.inp_modes = inp_modes \n",
    "        self.out_modes = out_modes \n",
    "        self.inp_size = np.prod(inp_modes)\n",
    "        self.out_size = np.prod(out_modes)\n",
    "        self.ranks = ranks\n",
    "        \n",
    "        self.W_cores = self.init_W_cores(inp_modes, out_modes)\n",
    "        self.b = torch.nn.Parameter(torch.ones(self.out_size))\n",
    "        \n",
    "    def init_W_cores(self, inp_modes, out_modes):\n",
    "        cores = torch.nn.ParameterList()\n",
    "        for k in range(self.d):\n",
    "            core = torch.randn(self.ranks[k], inp_modes[k], out_modes[k], self.ranks[k+1])\n",
    "            core *= 2 / (inp_modes[k] + out_modes[k])\n",
    "            cores.append(torch.nn.Parameter(core))\n",
    "        return cores\n",
    "            \n",
    "    def forward(self, inp):\n",
    "        W = TensorTrain(self.W_cores, self.inp_modes, self.out_modes, self.ranks)\n",
    "        out = inp * W + self.b\n",
    "        return out\n",
    "    \n",
    "    def backward(self, out):\n",
    "        grad = self.init_W_cores(inp_modes, out_modes)\n",
    "        return grad\n",
    "    \n",
    "    \n",
    "class My_Optimizer(torch.optim.Optimizer):\n",
    "    def __init__(self, optimizer):\n",
    "        self.optimizer = optimizer\n",
    "    \n",
    "#     def step(self):\n",
    "#         pass\n",
    "    \n",
    "    def __getattr__(self, attrname):\n",
    "        print(attrname)\n",
    "        return getattr(self.optimizer, attrname) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "        \n",
    "        self.net= nn.Sequential(\n",
    "            my_TT_linear([8, 8, 8], [4, 4, 4], [1, 2, 2, 1]), nn.ReLU(),\n",
    "            nn.Linear(64, 10), nn.Softmax(dim=1),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 1000\n",
    "inp_size = 512\n",
    "out_size = 10\n",
    "\n",
    "X = torch.rand(N, inp_size)\n",
    "W_true = torch.rand(inp_size, out_size) * 2 / (inp_size + out_size)\n",
    "y = torch.argmax(X @ W_true, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "\n",
    "trainset = TensorDataset(X, y)\n",
    "train_dataloader = torch.utils.data.DataLoader(trainset, batch_size=16, shuffle=True, num_workers=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(2.3556, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3503, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3556, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3340, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3193, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3523, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3386, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3487, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3413, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3320, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3386, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3367, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3385, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3290, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3271, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3462, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3114, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3233, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3239, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3315, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3435, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3449, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3220, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3336, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3336, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3229, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3366, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3365, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3293, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3256, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3424, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3008, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3213, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3242, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3141, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3341, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3372, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3232, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3304, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3212, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3211, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3061, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3208, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3080, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3129, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3196, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3177, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3169, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3084, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3123, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3119, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3097, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3066, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3155, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3029, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3180, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3180, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3111, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3187, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3074, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3078, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.2957, grad_fn=<NllLossBackward0>) torch.Size([16]) torch.Size([])\n",
      "tensor(2.3093, grad_fn=<NllLossBackward0>) torch.Size([8]) torch.Size([])\n"
     ]
    }
   ],
   "source": [
    "model = Net().to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",
    "\n",
    "train_loss, train_accuracy = train(model, train_dataloader, \n",
    "                                   criterion, optimizer, \n",
    "                                   n_epochs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_zero_grad_profile_name\n",
      "_zero_grad_profile_name\n",
      "param_groups\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "step() missing 1 required positional argument: 'closure'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-20-13694ab6cdb4>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0moptimizer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mMy_Optimizer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m train_loss, train_accuracy = train(model, train_dataloader, \n\u001b[0m\u001b[1;32m      8\u001b[0m                                    \u001b[0mcriterion\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m                                    n_epochs=5, show=True)\n",
      "\u001b[0;32m~/Desktop/Tensor_Train/lib/train.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(model, train_dataloader, criterion, optimizer, n_epochs, show, verbose)\u001b[0m\n\u001b[1;32m     56\u001b[0m             \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     57\u001b[0m             \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[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 58\u001b[0;31m             \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[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     59\u001b[0m             \u001b[0mloss_trace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\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     60\u001b[0m             \u001b[0my_pred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpredict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: step() missing 1 required positional argument: 'closure'"
     ]
    }
   ],
   "source": [
    "model = Net().to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",
    "optimizer = My_Optimizer(optimizer)\n",
    "\n",
    "train_loss, train_accuracy = train(model, train_dataloader, \n",
    "                                   criterion, optimizer, \n",
    "                                   n_epochs=5, show=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<bound method Optimizer.load_state_dict of My_Optimizer (\n",
       "Parameter Group 0\n",
       "    amsgrad: False\n",
       "    betas: (0.9, 0.999)\n",
       "    eps: 1e-08\n",
       "    lr: 0.0001\n",
       "    weight_decay: 0\n",
       ")>"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimizer.load_state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "base"
  },
  "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}