numpy_nn/test/module_tests.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "427c3428",
   "metadata": {},
   "outputs": [],
   "source": [
    "# move to project's root directory to make\n",
    "# numpy_nn and pytorch_nn packages accessable \n",
    "# %cd ../..\n",
    "\n",
    "\n",
    "# # Another possible solution is appending to the sys.path\n",
    "\n",
    "import sys\n",
    "import  os\n",
    "project_root = os.path.dirname(os.path.dirname(sys.path[0]))\n",
    "if project_root not in sys.path:\n",
    "    sys.path.append(project_root)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "59a56421",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Tuple, Callable\n",
    "\n",
    "import torch\n",
    "torch.use_deterministic_algorithms(True)\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import f1_score\n",
    "\n",
    "from test_layer import test_module_randomly\n",
    "\n",
    "from numpy_nn.modules.np_nn import (\n",
    "    FullyConnectedLayer,\n",
    "    ReLULayer,\n",
    "    SigmoidLayer,\n",
    "    ReLULayer,\n",
    "    AdamOptimizer,\n",
    "    CrossEntropyLoss,\n",
    "    LinearActivation,\n",
    "    Sequential,\n",
    "    Optimizer,\n",
    "    SoftMaxLayer,\n",
    "    GradientDescentOptimizer,\n",
    "    CrossEntropyLossWithSoftMax,\n",
    "    Conv2d,\n",
    "    Conv2dWithLoops,\n",
    "    Flatten,\n",
    "    MaxPool2d,\n",
    "    AdamOptimizer,\n",
    "    BatchNormalization2d,\n",
    "    ActivationLayer,\n",
    "    TrainableLayer,\n",
    ")\n",
    "\n",
    "from pytorch_nn.models.resnet import (\n",
    "    Bottleneck as Bottleneck_torch,\n",
    "    resnet101 as resnet101_torch\n",
    ")\n",
    "\n",
    "# from pytorch_nn.models.resnet_without_batchnorm import (\n",
    "#     Bottleneck as Bottleneck_torch_without_batchnorm,\n",
    "#     resnet101 as resnet101_torch_without_batchnorm\n",
    "# )\n",
    "\n",
    "from numpy_nn.models.resnet import Bottleneck, resnet101\n",
    "\n",
    "# from numpy_nn.models.resnet_without_batchnorm import (\n",
    "#     Bottleneck as Bottleneck_np_without_batchnorm,\n",
    "#     resnet101 as resnet101_np_without_batchnorm\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d8fb738f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import  os\n",
    "project_root = os.path.dirname(os.path.dirname(sys.path[0]))\n",
    "if project_root not in sys.path:\n",
    "    sys.path.append(project_root)\n",
    "\n",
    "\n",
    "\n",
    "from typing import Callable, Tuple, List\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from numpy_nn.modules.np_nn import (\n",
    "    FullyConnectedLayer,\n",
    "    BatchNormalization2d,\n",
    "    TrainableLayer,\n",
    "    Module,\n",
    ")\n",
    "\n",
    "\n",
    "def copy_trainable_parameters(my_module: Module, torch_module: torch.nn.Module) -> None:\n",
    "    if isinstance(my_module, FullyConnectedLayer):\n",
    "        my_module.weights = torch_module.weight.detach().numpy().T\n",
    "        my_module.bias = torch_module.bias.detach().numpy().reshape(-1, 1).T\n",
    "    elif isinstance(my_module, BatchNormalization2d):\n",
    "        n_channels = my_module.n_channels\n",
    "        my_module.gamma = torch_module.weight.detach().numpy().reshape(1, n_channels, 1, 1)\n",
    "        my_module.beta = torch_module.bias.detach().numpy().reshape(1, n_channels, 1, 1)\n",
    "        my_module.running_mean = torch_module.running_mean.detach().numpy().reshape(1, n_channels, 1, 1)\n",
    "        my_module.running_var = torch_module.running_var.detach().numpy().reshape(1, n_channels, 1, 1)\n",
    "    else:\n",
    "        my_module.weights = torch_module.weight.detach().numpy()\n",
    "        if my_module.bias is not None:\n",
    "            my_module.bias = torch_module.bias.detach().numpy()\n",
    "\n",
    "\n",
    "def test_module_randomly(my_module: Module,\n",
    "                torch_module: torch.nn.Module,\n",
    "                input_shape: Tuple[int, ...],\n",
    "                output_shape: Tuple[int, ...],\n",
    "                atol: float = 1e-5,\n",
    "                random_sampler: Callable = np.random.rand,\n",
    "                skip_parameter_copying: bool = False,\n",
    "                print_tensors: bool = False,\n",
    "                print_results: bool = False) -> None:\n",
    "    \"\"\"\n",
    "    Compares the output and gradients of a numpy layer and a torch layer\n",
    "\n",
    "    Args:\n",
    "        my_module: neural network layer implemented in numpy.\n",
    "        torch_module: neural network layer implemented in torch.\n",
    "        input_shape: shape of the input tensor.\n",
    "        output_shape: shape of the output tensor. It's used to generate\n",
    "            a random tensor representing partial derivative of the loss\n",
    "            function with respect to the output of the layer.\n",
    "        atol: absolute tolerance for comparing\n",
    "            numpy and torch tensors (used in np.allclose).\n",
    "        random_sampler: function that generates random tensors of the given shape.\n",
    "        skip_parameter_copying: if True, the weights and biases will be held intact.\n",
    "            By default, weights and biases are copied from torch_module to my_module.\n",
    "    \"\"\"    \n",
    "    # copy weights from torch_module to my_module\n",
    "    # if the numpy layer is trainable\n",
    "    if not skip_parameter_copying and isinstance(my_module, TrainableLayer):\n",
    "        copy_trainable_parameters(my_module, torch_module)\n",
    "\n",
    "    input_np = random_sampler(*input_shape).astype(np.float32)\n",
    "    input_torch = torch.from_numpy(input_np)\n",
    "    input_torch.requires_grad = True\n",
    "\n",
    "    output_np = my_module.forward(input_np)\n",
    "    output_torch = torch_module(input_torch)\n",
    "\n",
    "\n",
    "    if print_tensors:\n",
    "        print(\"my and torch outputs:\")\n",
    "        print(output_np.flatten(), output_torch.detach().numpy().flatten())\n",
    "    assert np.allclose(output_np, output_torch.detach().numpy(), atol=atol), \"Outputs are not equal\"\n",
    "    if print_results:\n",
    "        print(\"Outputs are equal\")\n",
    "    \n",
    "    # print(output_np.dtype)\n",
    "    # print(output_torch.dtype)\n",
    "    # print(output_torch.detach().numpy().dtype)\n",
    "\n",
    "    output_grad_np = random_sampler(*output_shape)\n",
    "    output_grad_torch = torch.from_numpy(output_grad_np)\n",
    "\n",
    "    input_grad_np = my_module.backward(output_grad_np)\n",
    "    output_torch.backward(output_grad_torch)\n",
    "    input_grad_torch = input_torch.grad.detach().numpy()\n",
    "\n",
    "    if print_tensors:\n",
    "        print(\"my and torch input gradients:\")\n",
    "        print(input_grad_np.flatten(), input_grad_torch.flatten())\n",
    "    assert np.allclose(input_grad_np, input_grad_torch, atol=atol), \"Input gradients are not equal\"\n",
    "    if print_results:\n",
    "        print(\"Input gradients are equal\")\n",
    "\n",
    "\n",
    "    if not isinstance(my_module, TrainableLayer):\n",
    "        return\n",
    "    \n",
    "\n",
    "    # compare weight and bias gradients\n",
    "    if isinstance(my_module, FullyConnectedLayer):\n",
    "        weight_grad_np = my_module.weights_gradient\n",
    "        weight_grad_torch = torch_module.weight.grad.detach().numpy().T\n",
    "        bias_grad_np = my_module.bias_gradient\n",
    "        bias_grad_torch = torch_module.bias.grad.detach().numpy().reshape(-1, 1).T\n",
    "    elif isinstance(my_module, BatchNormalization2d):\n",
    "        weight_grad_np = my_module.gamma_gradient.flatten()\n",
    "        weight_grad_torch = torch_module.weight.grad.detach().numpy()\n",
    "        bias_grad_np = my_module.beta_gradient.flatten()\n",
    "        bias_grad_torch = torch_module.bias.grad.detach().numpy()\n",
    "\n",
    "        if print_tensors:\n",
    "            print(\"my and torch running means:\")\n",
    "            print(my_module.running_mean.flatten(), torch_module.running_mean.detach().numpy().flatten())\n",
    "        running_mean_close = np.allclose(\n",
    "            my_module.running_mean.flatten(), torch_module.running_mean.detach().numpy().flatten(), atol=atol)\n",
    "        assert running_mean_close, \"Running mean is not equal\"\n",
    "        if print_results:\n",
    "            print(\"Running means are equal\")\n",
    "\n",
    "        if print_tensors:\n",
    "            print(\"my and torch running vars:\")\n",
    "            print(my_module.running_var.flatten(), torch_module.running_var.detach().numpy().flatten())\n",
    "        running_var_close = np.allclose(\n",
    "            my_module.running_var.flatten(), torch_module.running_var.detach().numpy().flatten(), atol=atol)\n",
    "        assert running_var_close, \"Running var is not equal\"\n",
    "        if print_results:\n",
    "            print(\"Running vars are equal\")\n",
    "            \n",
    "    else:\n",
    "        weight_grad_np = my_module.weights_gradient\n",
    "        weight_grad_torch = torch_module.weight.grad.detach().numpy()\n",
    "        if my_module.bias is not None:\n",
    "            bias_grad_np = my_module.bias_gradient\n",
    "            bias_grad_torch = torch_module.bias.grad.detach().numpy()\n",
    "    \n",
    "    weight_grads_close = np.allclose(weight_grad_np, weight_grad_torch, atol=atol)\n",
    "\n",
    "    if print_tensors:\n",
    "        print(\"my and torch weight gradients:\")\n",
    "        print(weight_grad_np.flatten(), weight_grad_torch.flatten())\n",
    "        \n",
    "    assert weight_grads_close, \"Weight gradients are not equal\"\n",
    "    if print_results:\n",
    "        print(\"Weight gradients are equal\")\n",
    "\n",
    "    if isinstance(my_module, BatchNormalization2d) or my_module.bias is not None:\n",
    "        if print_tensors:\n",
    "            print(\"my and torch bias gradients:\")\n",
    "            print(bias_grad_np.flatten(), bias_grad_torch.flatten())\n",
    "\n",
    "        assert np.allclose(bias_grad_np, bias_grad_torch, atol=atol), \"Bias gradients are not equal\"\n",
    "        if print_results:\n",
    "            print(\"Bias gradients are equal\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b38f73c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # reload a user's module test_layer\n",
    "# %load_ext autoreload\n",
    "# %autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "87c80e7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def batchnorm2d_iterative_test(n_channels: int, batch_size: int, height: int,\n",
    "                               width: int, n_iter: int, phase: str = \"train\",\n",
    "                               momentum: float = 0.1, atol: float = 1e-5,\n",
    "                               random_sampler: Callable = np.random.rand,\n",
    "                               print_tensors = False, print_results = False) -> None:\n",
    "            \n",
    "    # since each iteration changes the running mean, running variance, mean and variance\n",
    "    # this test runs several iterations and checks that the results are the same\n",
    "\n",
    "    # maybe there should be a custom test that calls\n",
    "    # forward multiple times and calls backward only once\n",
    "    \n",
    "    input_shape = output_shape = [batch_size, n_channels, height, width]\n",
    "\n",
    "    my_bn = BatchNormalization2d(n_channels, momentum=momentum)\n",
    "    torch_bn = torch.nn.BatchNorm2d(n_channels, momentum=momentum)\n",
    "    if phase == \"eval\":\n",
    "        my_bn.eval()\n",
    "        torch_bn.eval()\n",
    "    else:\n",
    "        my_bn.train()\n",
    "        torch_bn.train()\n",
    "\n",
    "    for i in range(n_iter):\n",
    "        print(\"Iteration:\", i)\n",
    "        # print(my_bn.training)\n",
    "        test_module_randomly(my_bn, torch_bn, input_shape, output_shape,\n",
    "                    atol=atol, random_sampler=random_sampler,\n",
    "                    skip_parameter_copying=bool(i),\n",
    "                    print_tensors=print_tensors,\n",
    "                    print_results=print_results)\n",
    "        # reset torch gradients\n",
    "        torch_bn.weight.grad = None\n",
    "        torch_bn.bias.grad = None\n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "57bd5af6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "atol = 1e-06, np.random.randn\n",
      "Iteration: 0\n",
      "\n",
      "Iteration: 1\n",
      "\n",
      "Iteration: 2\n",
      "\n",
      "Iteration: 3\n",
      "\n",
      "Iteration: 4\n",
      "\n",
      "Iteration: 5\n",
      "\n",
      "Iteration: 6\n",
      "\n",
      "Iteration: 7\n",
      "\n",
      "Iteration: 8\n",
      "\n",
      "Iteration: 9\n",
      "\n"
     ]
    }
   ],
   "source": [
    "atol = 1e-6\n",
    "print(f\"atol = {atol}, np.random.randn\")\n",
    "batchnorm2d_iterative_test(n_channels=5, batch_size=4, height=8, width=8,\n",
    "                           n_iter=10, phase=\"train\", momentum=0.1,\n",
    "                           atol=atol, random_sampler=np.random.randn, print_results=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "059674e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "atol = 0.0001, np.random.rand\n",
      "Iteration: 0\n",
      "\n",
      "Iteration: 1\n",
      "\n",
      "Iteration: 2\n",
      "\n",
      "Iteration: 3\n",
      "\n",
      "Iteration: 4\n",
      "\n",
      "Iteration: 5\n",
      "\n",
      "Iteration: 6\n",
      "\n",
      "Iteration: 7\n",
      "\n",
      "Iteration: 8\n",
      "\n",
      "Iteration: 9\n",
      "\n"
     ]
    }
   ],
   "source": [
    "atol = 1e-4\n",
    "print(f\"atol = {atol}, np.random.rand\")\n",
    "batchnorm2d_iterative_test(n_channels=5, batch_size=4, height=8, width=8,\n",
    "                           n_iter=10, phase=\"train\", momentum=0.8,\n",
    "                           atol=atol, random_sampler=np.random.rand, print_results=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "24933141",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "atol = 0.0001, np.random.rand\n",
      "Iteration: 0\n",
      "my and torch outputs:\n",
      "[-1.4301041  -1.082526   -0.19767104 ...  0.9150699   0.8313913\n",
      " -1.6011153 ] [-1.4301041  -1.082526   -0.19767106 ...  0.9150702   0.83139145\n",
      " -1.6011153 ]\n",
      "my and torch input gradients:\n",
      "[ 0.136774040017978  1.214768857994491  0.724841884986875 ...\n",
      " -0.610479597399505 -0.378000684800799 -1.016681103117026] [ 0.1367743   1.2147692   0.724842   ... -0.6104797  -0.37800083\n",
      " -1.0166808 ]\n",
      "my and torch running means:\n",
      "[0.3952297  0.38302088 0.39440167 0.3867925  0.3916015 ] [0.39522967 0.3830209  0.39440167 0.38679248 0.39160147]\n",
      "my and torch running vars:\n",
      "[0.2785678  0.2648916  0.26677704 0.26196226 0.27252778] [0.2785678  0.26489156 0.26677704 0.26196226 0.27252775]\n",
      "my and torch weight gradients:\n",
      "[-2.483329650028675  2.095407849137224 -2.46034050029717\n",
      "  2.691619001228386 -3.71991454658423 ] [-2.4833174  2.095394  -2.4603398  2.6916323 -3.7199028]\n",
      "my and torch bias gradients:\n",
      "[128.74741043319176 125.57231519081566 114.56492972465452\n",
      " 125.15025584738889 129.23468469521472] [128.7474   125.57231  114.564926 125.15025  129.23468 ]\n",
      "\n",
      "Iteration: 1\n",
      "my and torch outputs:\n",
      "[ 0.3457517  -0.73298854 -1.0806885  ...  1.1077915   0.19281165\n",
      " -1.6204898 ] [ 0.34575152 -0.73298866 -1.0806885  ...  1.1077918   0.19281185\n",
      " -1.6204896 ]\n",
      "my and torch input gradients:\n",
      "[ 0.840052785456636  1.681024713454835 -1.661235436599868 ...\n",
      "  1.484538696609036  0.270236300697097  1.163098096363204] [ 0.84005284  1.6810244  -1.6612357  ...  1.4845382   0.27023634\n",
      "  1.1630988 ]\n",
      "my and torch running means:\n",
      "[0.45220733 0.49917576 0.48977995 0.46549153 0.48210877] [0.45220733 0.49917576 0.48977995 0.4654915  0.4821087 ]\n",
      "my and torch running vars:\n",
      "[0.124127135 0.12364505  0.12116127  0.12119238  0.114460036] [0.124127135 0.123645045 0.12116125  0.12119239  0.114460036]\n",
      "my and torch weight gradients:\n",
      "[ 1.685880094062552 -3.936791963991086  7.324884295270801\n",
      " -0.229268255931551 -7.14409120340518 ] [ 1.6858667  -3.9367926   7.3248863  -0.22926833 -7.144062  ]\n",
      "my and torch bias gradients:\n",
      "[130.28944191375703 135.09011143653424 125.51532932647785\n",
      " 122.76045299231421 133.22212742133576] [130.28944 135.09012 125.51533 122.76046 133.22212]\n",
      "\n",
      "Iteration: 2\n",
      "my and torch outputs:\n",
      "[ 0.86728525 -1.2525281  -0.3090039  ...  1.6123682  -0.11933849\n",
      "  1.1666685 ] [ 0.86728525 -1.2525282  -0.30900395 ...  1.6123682  -0.11933851\n",
      "  1.1666685 ]\n",
      "my and torch input gradients:\n",
      "[-1.430807499938515 -0.304293609707524 -0.995636143295636 ...\n",
      " -0.783762347618578 -0.758927925738763 -1.476359301133926] [-1.4308074  -0.3042937  -0.99563617 ... -0.7837621  -0.75892794\n",
      " -1.476359  ]\n",
      "my and torch running means:\n",
      "[0.48370367 0.5147653  0.50263214 0.49730897 0.48205352] [0.4837037  0.5147653  0.5026322  0.4973089  0.48205352]\n",
      "my and torch running vars:\n",
      "[0.10026027 0.09197707 0.08994842 0.08759581 0.08787933] [0.10026027 0.09197707 0.08994843 0.08759582 0.08787933]\n",
      "my and torch weight gradients:\n",
      "[-4.678491169915187  2.223387676136488 -5.347108834551229\n",
      "  2.446429913252232 -2.38625213477397 ] [-4.6785026  2.2233884 -5.3471365  2.4464564 -2.3862638]\n",
      "my and torch bias gradients:\n",
      "[118.78381878435846 127.31572523528206 126.27130253508085\n",
      " 124.98528088616644 123.18785387367505] [118.78381 127.31572 126.2713  124.98528 123.18785]\n",
      "\n",
      "Iteration: 3\n",
      "my and torch outputs:\n",
      "[ 0.4701708  -0.36615634  0.16886832 ...  0.2692886   0.9150766\n",
      " -0.24463531] [ 0.47017097 -0.36615622  0.16886854 ...  0.26928866  0.9150766\n",
      " -0.24463534]\n",
      "my and torch input gradients:\n",
      "[ 1.029017429562337 -0.06917436298286  -0.047746987422307 ...\n",
      " -1.102606820763847 -1.138033488892933  0.567302311606553] [ 1.0290172   -0.06917428  -0.047747068 ... -1.1026068   -1.1380334\n",
      "  0.5673024  ]\n",
      "my and torch running means:\n",
      "[0.5091306  0.50381505 0.51026535 0.47424063 0.4675054 ] [0.50913054 0.50381505 0.5102654  0.4742406  0.46750537]\n",
      "my and torch running vars:\n",
      "[0.08716897 0.08400565 0.08518496 0.07876741 0.08327579] [0.08716897  0.08400564  0.085184954 0.07876741  0.08327578 ]\n",
      "my and torch weight gradients:\n",
      "[ 4.806328315972165  3.965155937105467 -5.918247095456863\n",
      " 10.023250102191966 -4.760757118373716] [ 4.806353   3.965156  -5.9182725 10.02325   -4.760756 ]\n",
      "my and torch bias gradients:\n",
      "[121.69418624758913 121.50618842661446 123.21829207983842\n",
      " 129.85770429303074 125.41166198569871] [121.69419 121.50619 123.21829 129.8577  125.41166]\n",
      "\n",
      "Iteration: 4\n",
      "my and torch outputs:\n",
      "[ 0.3123651  0.5838685  0.7852571 ... -1.1951635  0.745132   0.5189042] [ 0.31236517  0.5838686   0.7852572  ... -1.1951635   0.74513197\n",
      "  0.5189042 ]\n",
      "my and torch input gradients:\n",
      "[ 1.445655304321676 -0.474372390499244  0.107445394984742 ...\n",
      " -0.417721272486481  0.401912649836681  1.227752337764261] [ 1.4456555  -0.47437236  0.10744547 ... -0.4177213   0.40191254\n",
      "  1.2277523 ]\n",
      "my and torch running means:\n",
      "[0.4978351  0.49768093 0.50117815 0.49688813 0.5122338 ] [0.49783507 0.4976809  0.5011782  0.4968881  0.5122338 ]\n",
      "my and torch running vars:\n",
      "[0.08446134 0.08474696 0.08379264 0.08109796 0.0837315 ] [0.08446133  0.08474695  0.083792634 0.08109796  0.0837315  ]\n",
      "my and torch weight gradients:\n",
      "[-1.804141130421917 -2.495052060113101 -8.41020456600037\n",
      " -4.556111897941046 -4.18715183698759 ] [-1.8041408 -2.4950378 -8.410219  -4.5561104 -4.187151 ]\n",
      "my and torch bias gradients:\n",
      "[125.6244535222202  121.16767872773971 131.90143321054626\n",
      " 129.69679705853582 127.17118537181645] [125.62445 121.16768 131.90143 129.6968  127.17119]\n",
      "\n",
      "Iteration: 5\n",
      "my and torch outputs:\n",
      "[ 0.22456375  0.6147539   0.660237   ...  1.4729288   0.41555664\n",
      " -0.5711443 ] [ 0.22456372  0.61475384  0.66023695 ...  1.4729286   0.41555655\n",
      " -0.5711442 ]\n",
      "my and torch input gradients:\n",
      "[-1.571941715536342  1.522772330911733 -0.002280098127947 ...\n",
      " -1.157851876808562  0.339463472393153 -0.590522784677984] [-1.5719417     1.5227723    -0.0022800632 ... -1.1578519     0.33946335\n",
      " -0.5905228   ]\n",
      "my and torch running means:\n",
      "[0.5050174  0.52036446 0.52154875 0.4976419  0.4894219 ] [0.5050174  0.5203644  0.52154875 0.49764186 0.4894219 ]\n",
      "my and torch running vars:\n",
      "[0.079279654 0.08921964  0.078204095 0.07993992  0.086105764] [0.07927965  0.08921964  0.078204095 0.07993992  0.08610577 ]\n",
      "my and torch weight gradients:\n",
      "[ 0.023070493278048 -1.712801153762796 -0.497475748327792\n",
      " -0.125557127583595  2.432352596887311] [ 0.02307027 -1.7128017  -0.49747512 -0.12555654  2.4323523 ]\n",
      "my and torch bias gradients:\n",
      "[130.89094642580005 131.3328602421737  130.07454721858812\n",
      " 129.95597490747969 129.8700967275123 ] [130.89095 131.33286 130.07455 129.95598 129.8701 ]\n",
      "\n",
      "Iteration: 6\n",
      "my and torch outputs:\n",
      "[-1.2794167   0.51479477  0.77692753 ... -0.15602502  1.1202376\n",
      " -0.52973354] [-1.2794166   0.51479506  0.77692795 ... -0.15602505  1.1202376\n",
      " -0.52973354]\n",
      "my and torch input gradients:\n",
      "[-0.648736111517411 -0.83778349497227  -1.337774215233054 ...\n",
      "  0.995685915829544 -1.436236315968045  0.127105865232353] [-0.6487357  -0.8377837  -1.3377745  ...  0.99568594 -1.4362361\n",
      "  0.12710597]\n",
      "my and torch running means:\n",
      "[0.50589067 0.499444   0.51397884 0.5278444  0.5027191 ] [0.50589067 0.49944395 0.51397884 0.52784437 0.5027191 ]\n",
      "my and torch running vars:\n",
      "[0.087353684 0.08960599  0.08676637  0.080574036 0.083986074] [0.087353684 0.08960599  0.08676638  0.080574036 0.083986074]\n",
      "my and torch weight gradients:\n",
      "[-0.015450979146019  2.682995079433835  1.532254995312307\n",
      "  1.642005846049513  1.445104499847517] [-0.0154259205  2.6830072     1.5322552     1.6420058     1.4451059   ]\n",
      "my and torch bias gradients:\n",
      "[126.68701378126607 125.97642064368102 139.0777684622568\n",
      " 128.80666698509816 133.16739450050727] [126.68701  125.976425 139.07777  128.80667  133.16739 ]\n",
      "\n",
      "Iteration: 7\n",
      "my and torch outputs:\n",
      "[-0.3981201   1.596546   -0.90829724 ... -0.85748416  0.89456725\n",
      " -1.3922858 ] [-0.39812028  1.596546   -0.90829736 ... -0.85748416  0.89456725\n",
      " -1.3922858 ]\n",
      "my and torch input gradients:\n",
      "[-0.849814154407687 -0.583633336609537 -0.937710459923952 ...\n",
      "  0.360539856302955  0.634777554405773  0.96691812731229 ] [-0.8498143  -0.583633   -0.93771064 ...  0.36053994  0.6347776\n",
      "  0.9669182 ]\n",
      "my and torch running means:\n",
      "[0.4998647  0.5080074  0.5127528  0.51365876 0.50729936] [0.4998647  0.50800735 0.5127528  0.5136587  0.50729936]\n",
      "my and torch running vars:\n",
      "[0.08065316  0.09139621  0.087348975 0.084453605 0.0804811  ] [0.08065315  0.091396205 0.08734897  0.084453605 0.08048109 ]\n",
      "my and torch weight gradients:\n",
      "[ 2.906337850712552  0.821341293543365 -3.953726534200371\n",
      " -3.792370132985228  0.233259888859514] [ 2.9063246   0.82134145 -3.9537256  -3.7923708   0.23325926]\n",
      "my and torch bias gradients:\n",
      "[132.0598717262929  119.9673581248816  130.64518448750698\n",
      " 126.39592052064557 131.87435382649338] [132.05988 119.96736 130.64519 126.39592 131.87434]\n",
      "\n",
      "Iteration: 8\n",
      "my and torch outputs:\n",
      "[-1.399553     0.100612774  1.0429701   ... -0.016142836 -0.3282107\n",
      "  0.3847755  ] [-1.3995528    0.10061288   1.0429702   ... -0.016142845 -0.3282107\n",
      "  0.38477552 ]\n",
      "my and torch input gradients:\n",
      "[-0.868524861179574  0.207994442464159  1.123850731276564 ...\n",
      " -0.517580304364597  0.806194100759699  0.561846425312011] [-0.8685243   0.2079943   1.1238502  ... -0.51758033  0.8061941\n",
      "  0.5618464 ]\n",
      "my and torch running means:\n",
      "[0.51837677 0.49841926 0.4963262  0.4830934  0.5063752 ] [0.5183767  0.49841923 0.4963262  0.48309338 0.5063752 ]\n",
      "my and torch running vars:\n",
      "[0.07548793 0.08064313 0.08268702 0.08804274 0.08641509] [0.07548792  0.080643125 0.08268702  0.08804274  0.08641509 ]\n",
      "my and torch weight gradients:\n",
      "[ 4.844621703464032 -1.152577818768132 -0.247894569615031\n",
      " -1.224259745061717  0.168297396709083] [ 4.8446484  -1.1525785  -0.2478935  -1.2242595   0.16829738]\n",
      "my and torch bias gradients:\n",
      "[124.7424137411785  125.03464160317864 128.8087492843802\n",
      " 129.1459334163105  122.93163970585486] [124.74242 125.03464 128.80875 129.14594 122.93164]\n",
      "\n",
      "Iteration: 9\n",
      "my and torch outputs:\n",
      "[-0.22394337 -0.40622693  0.06664393 ... -1.5528905  -1.0567306\n",
      "  1.1780813 ] [-0.22394335  -0.406227     0.066643834 ... -1.5528904   -1.0567305\n",
      "  1.1780814  ]\n",
      "my and torch input gradients:\n",
      "[-0.330272820619513 -0.050120403398395  1.029178495599392 ...\n",
      "  0.474690186997089 -1.730510460898985 -1.296645496521766] [-0.33027282  -0.050120402  1.0291785   ...  0.47469047  -1.7305101\n",
      " -1.2966455  ]\n",
      "my and torch running means:\n",
      "[0.5084295  0.48991776 0.51061654 0.5169275  0.4989316 ] [0.5084295  0.48991776 0.5106166  0.5169275  0.49893156]\n",
      "my and torch running vars:\n",
      "[0.081693515 0.08196727  0.07641114  0.082329646 0.08628556 ] [0.08169351  0.08196727  0.07641114  0.08232965  0.086285554]\n",
      "my and torch weight gradients:\n",
      "[ 8.564006676488052 -4.106699491764033 -2.833182343023842\n",
      "  4.751751350556482 -3.598384943928523] [ 8.564007  -4.106698  -2.8332107  4.751752  -3.5983713]\n",
      "my and torch bias gradients:\n",
      "[125.28902204391022 136.65398744807004 129.54222361985848\n",
      " 137.1685480242214  127.45799851769834] [125.289024 136.65398  129.54222  137.16855  127.45799 ]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "atol = 1e-4\n",
    "print(f\"atol = {atol}, np.random.rand\")\n",
    "batchnorm2d_iterative_test(n_channels=5, batch_size=4, height=8, width=8,\n",
    "                           n_iter=10, phase=\"train\", momentum=0.8,\n",
    "                           atol=atol, random_sampler=np.random.rand, print_results=False, print_tensors=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "6a0bd846",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "atol = 1e-05, np.random.rand\n",
      "Iteration: 0\n",
      "Outputs are equal\n",
      "Input gradients are equal\n",
      "Running means are equal\n",
      "Running vars are equal\n"
     ]
    },
    {
     "ename": "AssertionError",
     "evalue": "Weight gradients are not equal",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[1;32mc:\\Users\\SystemPoint\\Documents\\GitHub\\DNN_lab_1\\numpy_nn\\test\\module_tests.ipynb Cell 10\u001b[0m in \u001b[0;36m<cell line: 3>\u001b[1;34m()\u001b[0m\n\u001b[0;32m      <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=0'>1</a>\u001b[0m atol \u001b[39m=\u001b[39m \u001b[39m1e-5\u001b[39m\n\u001b[0;32m      <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=1'>2</a>\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39matol = \u001b[39m\u001b[39m{\u001b[39;00matol\u001b[39m}\u001b[39;00m\u001b[39m, np.random.rand\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m----> <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=2'>3</a>\u001b[0m batchnorm2d_iterative_test(n_channels\u001b[39m=\u001b[39;49m\u001b[39m5\u001b[39;49m, batch_size\u001b[39m=\u001b[39;49m\u001b[39m4\u001b[39;49m, height\u001b[39m=\u001b[39;49m\u001b[39m8\u001b[39;49m, width\u001b[39m=\u001b[39;49m\u001b[39m8\u001b[39;49m,\n\u001b[0;32m      <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=3'>4</a>\u001b[0m                            n_iter\u001b[39m=\u001b[39;49m\u001b[39m10\u001b[39;49m, phase\u001b[39m=\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39mtrain\u001b[39;49m\u001b[39m\"\u001b[39;49m, momentum\u001b[39m=\u001b[39;49m\u001b[39m0.8\u001b[39;49m,\n\u001b[0;32m      <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=4'>5</a>\u001b[0m                            atol\u001b[39m=\u001b[39;49matol, random_sampler\u001b[39m=\u001b[39;49mnp\u001b[39m.\u001b[39;49mrandom\u001b[39m.\u001b[39;49mrand, print_results\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m)\n",
      "\u001b[1;32mc:\\Users\\SystemPoint\\Documents\\GitHub\\DNN_lab_1\\numpy_nn\\test\\module_tests.ipynb Cell 10\u001b[0m in \u001b[0;36mbatchnorm2d_iterative_test\u001b[1;34m(n_channels, batch_size, height, width, n_iter, phase, momentum, atol, random_sampler, print_tensors, print_results)\u001b[0m\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=124'>125</a>\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mIteration:\u001b[39m\u001b[39m\"\u001b[39m, i)\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=125'>126</a>\u001b[0m \u001b[39m# print(my_bn.training)\u001b[39;00m\n\u001b[1;32m--> <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=126'>127</a>\u001b[0m test_module(my_bn, torch_bn, input_shape, output_shape,\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=127'>128</a>\u001b[0m             atol\u001b[39m=\u001b[39;49matol, random_sampler\u001b[39m=\u001b[39;49mrandom_sampler,\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=128'>129</a>\u001b[0m             skip_parameter_copying\u001b[39m=\u001b[39;49m\u001b[39mbool\u001b[39;49m(i),\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=129'>130</a>\u001b[0m             print_tensors\u001b[39m=\u001b[39;49mprint_tensors,\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=130'>131</a>\u001b[0m             print_results\u001b[39m=\u001b[39;49mprint_results)\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=131'>132</a>\u001b[0m \u001b[39m# reset torch gradients\u001b[39;00m\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=132'>133</a>\u001b[0m torch_bn\u001b[39m.\u001b[39mweight\u001b[39m.\u001b[39mgrad \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
      "\u001b[1;32mc:\\Users\\SystemPoint\\Documents\\GitHub\\DNN_lab_1\\numpy_nn\\test\\module_tests.ipynb Cell 10\u001b[0m in \u001b[0;36mtest_module\u001b[1;34m(my_module, torch_module, input_shape, output_shape, atol, random_sampler, skip_parameter_copying, print_tensors, print_results)\u001b[0m\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=145'>146</a>\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mmy and torch weight gradients:\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=146'>147</a>\u001b[0m     \u001b[39mprint\u001b[39m(weight_grad_np\u001b[39m.\u001b[39mflatten(), weight_grad_torch\u001b[39m.\u001b[39mflatten())\n\u001b[1;32m--> <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=148'>149</a>\u001b[0m \u001b[39massert\u001b[39;00m weight_grads_close, \u001b[39m\"\u001b[39m\u001b[39mWeight gradients are not equal\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=149'>150</a>\u001b[0m \u001b[39mif\u001b[39;00m print_results:\n\u001b[0;32m    <a href='vscode-notebook-cell:/c%3A/Users/SystemPoint/Documents/GitHub/DNN_lab_1/numpy_nn/test/module_tests.ipynb#W6sZmlsZQ%3D%3D?line=150'>151</a>\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mWeight gradients are equal\u001b[39m\u001b[39m\"\u001b[39m)\n",
      "\u001b[1;31mAssertionError\u001b[0m: Weight gradients are not equal"
     ]
    }
   ],
   "source": [
    "atol = 1e-5\n",
    "print(f\"atol = {atol}, np.random.rand\")\n",
    "batchnorm2d_iterative_test(n_channels=5, batch_size=4, height=8, width=8,\n",
    "                           n_iter=10, phase=\"train\", momentum=0.8,\n",
    "                           atol=atol, random_sampler=np.random.rand, print_results=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba7f546",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss_val all close: True\n",
      "loss gradients all close: True\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "CrossEntropyLoss test random.rand\n",
    "\"\"\"\n",
    "\n",
    "def one_hot(y: np.ndarray, n_classes: int):\n",
    "    encoded = np.zeros((y.size, n_classes))\n",
    "    encoded[np.arange(y.size), y] = 1\n",
    "    return encoded\n",
    "\n",
    "\n",
    "batch_size = 5\n",
    "n_classes = 3\n",
    "pred = np.random.rand(batch_size, n_classes).astype(np.float32)\n",
    "true = one_hot(np.random.randint(0, n_classes, batch_size), n_classes)\n",
    "pred_torch = torch.from_numpy(pred).float()\n",
    "true_torch = torch.from_numpy(true).float()\n",
    "pred_torch.requires_grad = True\n",
    "\n",
    "torch_loss  = torch.nn.CrossEntropyLoss()\n",
    "torch_loss_val = torch_loss(pred_torch, true_torch)\n",
    "torch_loss_val.backward()\n",
    "\n",
    "my_loss = CrossEntropyLossWithSoftMax()\n",
    "my_loss_val = my_loss.forward(pred, true)\n",
    "my_loss.backward()\n",
    "\n",
    "print(\"loss_val all close:\", np.allclose(my_loss_val, torch_loss_val.detach().numpy()))\n",
    "print(\"loss gradients all close:\", np.allclose(my_loss.backward(), pred_torch.grad))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a0aaf69",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss_val all close: True\n",
      "loss gradients all close: True\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "CrossEntropyLoss test random.randn\n",
    "\"\"\"\n",
    "\n",
    "def one_hot(y: np.ndarray, n_classes: int):\n",
    "    encoded = np.zeros((y.size, n_classes))\n",
    "    encoded[np.arange(y.size), y] = 1\n",
    "    return encoded\n",
    "\n",
    "\n",
    "batch_size = 5\n",
    "n_classes = 3\n",
    "pred = np.random.randn(batch_size, n_classes).astype(np.float32)\n",
    "true = one_hot(np.random.randint(0, n_classes, batch_size), n_classes)\n",
    "pred_torch = torch.from_numpy(pred).float()\n",
    "true_torch = torch.from_numpy(true).float()\n",
    "pred_torch.requires_grad = True\n",
    "\n",
    "torch_loss  = torch.nn.CrossEntropyLoss()\n",
    "torch_loss_val = torch_loss(pred_torch, true_torch)\n",
    "torch_loss_val.backward()\n",
    "\n",
    "my_loss = CrossEntropyLossWithSoftMax()\n",
    "my_loss_val = my_loss.forward(pred, true)\n",
    "my_loss.backward()\n",
    "\n",
    "print(\"loss_val all close:\", np.allclose(my_loss_val, torch_loss_val.detach().numpy()))\n",
    "print(\"loss gradients all close:\", np.allclose(my_loss.backward(), pred_torch.grad))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "9b5717f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1mafaik the tests are passed if atol = 1e-4 or smaller regardless the sampler. If atol = 1e-6 unifrom distribution leads to problems\u001b[0m\n",
      "\n",
      "sampler = rand\n",
      "stride = 1\n",
      "output all close: True\n",
      "input gradients all close: False\n",
      "conv1 weights gradients all close: True\n",
      "conv2 weights gradients all close: False\n",
      "conv3 weights gradients all close: False\n",
      "bn1 gamma gradients all close: False\n",
      "bn1 beta gradients all close: True\n",
      "bn2 gamma gradients all close: True\n",
      "bn2 beta gradients all close: True\n",
      "bn3 gamma gradients all close: True\n",
      "bn3 beta gradients all close: True\n",
      "\n",
      "stride = 2\n",
      "output all close: True\n",
      "input gradients all close: False\n",
      "conv1 weights gradients all close: True\n",
      "conv2 weights gradients all close: False\n",
      "conv3 weights gradients all close: False\n",
      "conv_to_match_dimensions weights gradients all close: False\n",
      "bn_for_residual gamma gradients all close: True\n",
      "bn_for_residual beta gradients all close: True\n",
      "bn1 gamma gradients all close: True\n",
      "bn1 beta gradients all close: True\n",
      "bn2 gamma gradients all close: True\n",
      "bn2 beta gradients all close: True\n",
      "bn3 gamma gradients all close: True\n",
      "bn3 beta gradients all close: True\n",
      "\n",
      "sampler = randn\n",
      "stride = 1\n",
      "output all close: True\n",
      "input gradients all close: True\n",
      "conv1 weights gradients all close: True\n",
      "conv2 weights gradients all close: True\n",
      "conv3 weights gradients all close: True\n",
      "bn1 gamma gradients all close: True\n",
      "bn1 beta gradients all close: True\n",
      "bn2 gamma gradients all close: True\n",
      "bn2 beta gradients all close: False\n",
      "bn3 gamma gradients all close: True\n",
      "bn3 beta gradients all close: True\n",
      "\n",
      "stride = 2\n",
      "output all close: True\n",
      "input gradients all close: True\n",
      "conv1 weights gradients all close: True\n",
      "conv2 weights gradients all close: True\n",
      "conv3 weights gradients all close: True\n",
      "conv_to_match_dimensions weights gradients all close: True\n",
      "bn_for_residual gamma gradients all close: True\n",
      "bn_for_residual beta gradients all close: True\n",
      "bn1 gamma gradients all close: True\n",
      "bn1 beta gradients all close: True\n",
      "bn2 gamma gradients all close: True\n",
      "bn2 beta gradients all close: True\n",
      "bn3 gamma gradients all close: True\n",
      "bn3 beta gradients all close: True\n",
      "\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "BottleNeckLayer test\n",
    "\"\"\"\n",
    "\n",
    "batch_size = 5\n",
    "in_channels = 8\n",
    "bottleneck_depth = 2\n",
    "width = 6\n",
    "height = 6\n",
    "\n",
    "expansion_factor = 4\n",
    "n_output_channels = bottleneck_depth * expansion_factor\n",
    "\n",
    "momentum = 0.1\n",
    "\n",
    "print(\"\\n\" + \"\\033[1m\" + \"afaik the tests are passed if atol = 1e-4 \"\\\n",
    "      \"or smaller regardless the sampler. If atol = 1e-6 \"\n",
    "      \"unifrom distribution leads to problems\" + \"\\033[0m\" + \"\\n\")\n",
    "\n",
    "for sampler in (np.random.rand, np.random.randn):\n",
    "    print(f\"sampler = {sampler.__name__}\")\n",
    "    for stride_for_downsampling in (1, 2):  # Checking both cases: no downsampling and downsampling\n",
    "        print(f\"stride = {stride_for_downsampling}\")\n",
    "        input_data = sampler(batch_size, in_channels, width, height).astype(np.float32)\n",
    "        input_data_torch = torch.from_numpy(input_data).float()\n",
    "        input_data_torch.requires_grad = True\n",
    "\n",
    "        output_width = width // stride_for_downsampling\n",
    "        output_height = height // stride_for_downsampling\n",
    "        output_gradient = sampler(batch_size, n_output_channels, output_width, output_height).astype(np.float32)\n",
    "\n",
    "        torch_bottleneck = Bottleneck_torch(in_channels, bottleneck_depth, stride_for_downsampling)\n",
    "        my_bottleneck = Bottleneck(in_channels, bottleneck_depth, stride_for_downsampling)\n",
    "\n",
    "        conv_layer_pairs = [\n",
    "            (my_bottleneck.conv1, torch_bottleneck.conv1),\n",
    "            (my_bottleneck.conv2, torch_bottleneck.conv2),\n",
    "            (my_bottleneck.conv3, torch_bottleneck.conv3)]\n",
    "\n",
    "        for my_conv, torch_conv in conv_layer_pairs:\n",
    "            my_conv.weights = torch_conv.weight.detach().numpy() #.reshape(my_conv.weights.shape)\n",
    "        \n",
    "        bn_pairs = [\n",
    "            (my_bottleneck.bn1, torch_bottleneck.bn1),\n",
    "            (my_bottleneck.bn2, torch_bottleneck.bn2),\n",
    "            (my_bottleneck.bn3, torch_bottleneck.bn3)]\n",
    "        \n",
    "        for my_bn, torch_bn in bn_pairs:\n",
    "            my_bn.gamma = torch_bn.weight.detach().numpy().reshape(my_bn.gamma.shape)\n",
    "            my_bn.beta = torch_bn.bias.detach().numpy().reshape(my_bn.beta.shape)\n",
    "            my_bn.running_mean = torch_bn.running_mean.detach().numpy().reshape(my_bn.running_mean.shape)\n",
    "            my_bn.running_var = torch_bn.running_var.detach().numpy().reshape(my_bn.running_var.shape)\n",
    "            my_bn.momentum = torch_bn.momentum = momentum\n",
    "        \n",
    "\n",
    "        if my_bottleneck.conv_to_match_dimensions:\n",
    "            my_bottleneck.conv_to_match_dimensions.weights = torch_bottleneck.conv_to_match_dimensions.weight.detach().numpy()\n",
    "            my_bottleneck.bn_for_residual.gamma = torch_bottleneck.bn_for_residual.weight.detach().numpy().reshape(my_bottleneck.bn_for_residual.gamma.shape)\n",
    "            my_bottleneck.bn_for_residual.beta = torch_bottleneck.bn_for_residual.bias.detach().numpy().reshape(my_bottleneck.bn_for_residual.beta.shape)\n",
    "            my_bottleneck.bn_for_residual.running_mean = torch_bottleneck.bn_for_residual.running_mean.detach().numpy().reshape(my_bottleneck.bn_for_residual.running_mean.shape)\n",
    "            my_bottleneck.bn_for_residual.running_var = torch_bottleneck.bn_for_residual.running_var.detach().numpy().reshape(my_bottleneck.bn_for_residual.running_var.shape)\n",
    "            torch_bottleneck.bn_for_residual.momentum = my_bottleneck.bn_for_residual.momentum = momentum\n",
    "        \n",
    "        my_bottleneck.train()\n",
    "        torch_bottleneck.train()\n",
    "\n",
    "        my_out = my_bottleneck.forward(input_data)\n",
    "        torch_out = torch_bottleneck(input_data_torch)\n",
    "\n",
    "        torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "        torch_input_g = input_data_torch.grad.detach().numpy()\n",
    "\n",
    "        my_input_g = my_bottleneck.backward(output_gradient)\n",
    "\n",
    "        atol = 1e-6\n",
    "        print(\"output all close:\", np.allclose(my_out, torch_out.detach().numpy(), atol=atol))\n",
    "        print(\"input gradients all close:\", np.allclose(my_input_g, torch_input_g, atol=atol))\n",
    "        print(\"conv1 weights gradients all close:\", np.allclose(my_bottleneck.conv1.weights_gradient, torch_bottleneck.conv1.weight.grad.detach().numpy(), atol=atol))\n",
    "        print(\"conv2 weights gradients all close:\", np.allclose(my_bottleneck.conv2.weights_gradient, torch_bottleneck.conv2.weight.grad.detach().numpy(), atol=atol))\n",
    "        print(\"conv3 weights gradients all close:\", np.allclose(my_bottleneck.conv3.weights_gradient, torch_bottleneck.conv3.weight.grad.detach().numpy(), atol=atol))\n",
    "        if my_bottleneck.conv_to_match_dimensions:\n",
    "            print(\"conv_to_match_dimensions weights gradients all close:\", np.allclose(my_bottleneck.conv_to_match_dimensions.weights_gradient, torch_bottleneck.conv_to_match_dimensions.weight.grad.detach().numpy(), atol=atol))\n",
    "            print(\"bn_for_residual gamma gradients all close:\", np.allclose(my_bottleneck.bn_for_residual.gamma_gradient.flatten(), torch_bottleneck.bn_for_residual.weight.grad.detach().numpy(), atol=atol))\n",
    "            print(\"bn_for_residual beta gradients all close:\", np.allclose(my_bottleneck.bn_for_residual.beta_gradient.flatten(), torch_bottleneck.bn_for_residual.bias.grad.detach().numpy(), atol=atol))\n",
    "        \n",
    "\n",
    "        print(\"bn1 gamma gradients all close:\", np.allclose(my_bottleneck.bn1.gamma_gradient.flatten(), torch_bottleneck.bn1.weight.grad.detach().numpy(), atol=atol))\n",
    "        print(\"bn1 beta gradients all close:\", np.allclose(my_bottleneck.bn1.beta_gradient.flatten(), torch_bottleneck.bn1.bias.grad.detach().numpy(), atol=atol))\n",
    "        print(\"bn2 gamma gradients all close:\", np.allclose(my_bottleneck.bn2.gamma_gradient.flatten(), torch_bottleneck.bn2.weight.grad.detach().numpy(), atol=atol))\n",
    "        print(\"bn2 beta gradients all close:\", np.allclose(my_bottleneck.bn2.beta_gradient.flatten(), torch_bottleneck.bn2.bias.grad.detach().numpy(), atol=atol))\n",
    "        print(\"bn3 gamma gradients all close:\", np.allclose(my_bottleneck.bn3.gamma_gradient.flatten(), torch_bottleneck.bn3.weight.grad.detach().numpy(), atol=atol))  \n",
    "        print(\"bn3 beta gradients all close:\", np.allclose(my_bottleneck.bn3.beta_gradient.flatten(), torch_bottleneck.bn3.bias.grad.detach().numpy(), atol=atol))  \n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09b68d93",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c4f96846",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output all close: False\n",
      "input gradients all close: False\n",
      "fc weights gradients all close: False\n",
      "fc bias gradients all close: True\n",
      "It's ok that gradients don't match in eval mode. Batch norm's backward is different in train and eval mode.\n",
      "I didn't implement eval backward since it won't ever be used in training the network\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "resnet 101 test\n",
    "\"\"\"\n",
    "\n",
    "batch_size = 10\n",
    "height = width = 32\n",
    "n_channels = 1\n",
    "\n",
    "input_data = np.random.rand(batch_size, n_channels, height, width).astype(np.float32)\n",
    "input_data_torch = torch.from_numpy(input_data).float()\n",
    "input_data_torch.requires_grad = True\n",
    "output_gradient = np.random.rand(batch_size, 10).astype(np.float32)\n",
    "\n",
    "torch_resnet = resnet101_torch(10, 1)\n",
    "my_resnet = resnet101(10, 1)\n",
    "torch_resnet.train()\n",
    "my_resnet.train()\n",
    "\n",
    "my_resnet.clone_weights_from_torch(torch_resnet)\n",
    "\n",
    "my_out = my_resnet.forward(input_data)\n",
    "\n",
    "torch_out = torch_resnet(input_data_torch)\n",
    "torch_out_np = torch_out.detach().numpy()\n",
    "\n",
    "torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "torch_input_g = input_data_torch.grad.detach().numpy()\n",
    "\n",
    "my_input_g = my_resnet.backward(output_gradient)\n",
    "\n",
    "atol=1e-4\n",
    "\n",
    "print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "print(\"input gradients all close:\", np.allclose(my_input_g.flatten(), torch_input_g.flatten(), atol=atol))\n",
    "\n",
    "print(\"fc weights gradients all close:\", np.allclose(my_resnet.fc.weights_gradient, torch_resnet.fc.weight.grad.detach().numpy().T, atol=atol))\n",
    "print(\"fc bias gradients all close:\", np.allclose(my_resnet.fc.bias_gradient, torch_resnet.fc.bias.grad.detach().numpy(), atol=atol))\n",
    "\n",
    "print(\"It's ok that gradients don't match in eval mode. \"\\\n",
    "      \"Batch norm's backward is different in train and eval mode.\\n\"\\\n",
    "      \"I didn't implement eval backward since it won't ever be used in training the network\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "bf8d0ebb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([-49.13575793,  75.4237405 ,  28.05470836, ...,  73.2363413 ,\n",
       "        -36.06450964,  57.12665886]),\n",
       " array([-44.48268 ,  64.007614,  27.99699 , ...,  74.460556, -35.49483 ,\n",
       "         52.25823 ], dtype=float32))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_input_g.flatten(), torch_input_g.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "44cf20de",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output all close: True\n",
      "input gradients all close: False\n",
      "fc weights gradients all close: True\n",
      "fc bias gradients all close: True\n"
     ]
    }
   ],
   "source": [
    "atol=1e-2\n",
    "\n",
    "print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "print(\"input gradients all close:\", np.allclose(my_input_g, torch_input_g, atol=atol))\n",
    "\n",
    "print(\"fc weights gradients all close:\", np.allclose(my_resnet.fc.weights_gradient, torch_resnet.fc.weight.grad.detach().numpy().T, atol=atol))\n",
    "print(\"fc bias gradients all close:\", np.allclose(my_resnet.fc.bias_gradient, torch_resnet.fc.bias.grad.detach().numpy(), atol=atol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "822c960b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output all close: True\n",
      "input gradients all close: True\n",
      "fc weights gradients all close: True\n",
      "fc bias gradients all close: True\n"
     ]
    }
   ],
   "source": [
    "atol=1e+2\n",
    "\n",
    "print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "print(\"input gradients all close:\", np.allclose(my_input_g, torch_input_g, atol=atol))\n",
    "\n",
    "print(\"fc weights gradients all close:\", np.allclose(my_resnet.fc.weights_gradient, torch_resnet.fc.weight.grad.detach().numpy().T, atol=atol))\n",
    "print(\"fc bias gradients all close:\", np.allclose(my_resnet.fc.bias_gradient, torch_resnet.fc.bias.grad.detach().numpy(), atol=atol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "511f0e72",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output all close: True\n",
      "input gradients all close: True\n",
      "fc weights gradients all close: True\n",
      "fc bias gradients all close: True\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "resnet 101 without batchnormtest\n",
    "\"\"\"\n",
    "\n",
    "batch_size = 10\n",
    "height = width = 32\n",
    "n_channels = 1\n",
    "\n",
    "input_data = np.random.rand(batch_size, n_channels, height, width).astype(np.float32)\n",
    "input_data_torch = torch.from_numpy(input_data).float()\n",
    "input_data_torch.requires_grad = True\n",
    "output_gradient = np.random.rand(batch_size, 10).astype(np.float32)\n",
    "\n",
    "torch_resnet = resnet101_torch_without_batchnorm(10, 1)\n",
    "torch_resnet.train()\n",
    "my_resnet = resnet101_np_without_batchnorm(10, 1)\n",
    "my_resnet.train()\n",
    "\n",
    "my_resnet.clone_weights_from_torch(torch_resnet)\n",
    "\n",
    "my_out = my_resnet.forward(input_data)\n",
    "\n",
    "torch_out = torch_resnet(input_data_torch)\n",
    "torch_out_np = torch_out.detach().numpy()\n",
    "\n",
    "torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "torch_input_g = input_data_torch.grad.detach().numpy()\n",
    "\n",
    "my_input_g = my_resnet.backward(output_gradient)\n",
    "\n",
    "atol=1e-6\n",
    "\n",
    "print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "print(\"input gradients all close:\", np.allclose(my_input_g, torch_input_g, atol=atol))\n",
    "\n",
    "print(\"fc weights gradients all close:\", np.allclose(my_resnet.fc.weights_gradient, torch_resnet.fc.weight.grad.detach().numpy().T, atol=atol))\n",
    "print(\"fc bias gradients all close:\", np.allclose(my_resnet.fc.bias_gradient, torch_resnet.fc.bias.grad.detach().numpy(), atol=atol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "14fb02cd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output all close: True\n",
      "input gradients all close: True\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "conv1 test\n",
    "\"\"\"\n",
    "\n",
    "batch_size = 10\n",
    "height = width = 32\n",
    "n_channels = 1\n",
    "\n",
    "input_data = np.random.rand(batch_size, n_channels, height, width).astype(np.float32)\n",
    "input_data_torch = torch.from_numpy(input_data).float()\n",
    "input_data_torch.requires_grad = True\n",
    "output_gradient = np.random.rand(batch_size, 64, 16, 16).astype(np.float32)\n",
    "\n",
    "torch_resnet = resnet101_torch(10, 1)\n",
    "torch_resnet.eval()\n",
    "my_resnet = resnet101(10, 1)\n",
    "my_resnet.eval()\n",
    "\n",
    "my_resnet.clone_weights_from_torch(torch_resnet)\n",
    "\n",
    "my_out = my_resnet.conv1.forward(input_data)\n",
    "\n",
    "torch_out = torch_resnet.conv1(input_data_torch)\n",
    "torch_out_np = torch_out.detach().numpy()\n",
    "\n",
    "torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "torch_input_g = input_data_torch.grad.detach().numpy()\n",
    "\n",
    "my_input_g = my_resnet.conv1.backward(output_gradient)\n",
    "\n",
    "atol=1e-6\n",
    "\n",
    "print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "print(\"input gradients all close:\", np.allclose(my_input_g, torch_input_g, atol=atol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "1356b5fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "phase: train\n",
      "output all close: True\n",
      "input gradients all close: True\n",
      "\n",
      "phase: eval\n",
      "output all close: True\n",
      "input gradients all close: False\n",
      "\n",
      "It's ok that gradients don't match in eval mode. Batch norm's backward is different in train and eval mode.\n",
      "I didn't implement eval backward since it won't ever be used in training the network\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "bn1 test\n",
    "\"\"\"\n",
    "\n",
    "batch_size = 10\n",
    "height = width = 16\n",
    "n_channels = 64\n",
    "\n",
    "n_classes = 10\n",
    "\n",
    "for phase in ['train', 'eval']:\n",
    "    print(f\"phase: {phase}\")\n",
    "\n",
    "    # ! Has been noticed that moving 4 lines below outside the loop leads to not passing tests\n",
    "    input_data = np.random.rand(batch_size, n_channels, height, width).astype(np.float32)\n",
    "    input_data_torch = torch.from_numpy(input_data).float()\n",
    "    input_data_torch.requires_grad = True\n",
    "    output_gradient = np.random.rand(batch_size, n_channels, height, width).astype(np.float32)\n",
    "\n",
    "    torch_resnet = resnet101_torch(n_classes, n_channels)\n",
    "    my_resnet = resnet101(n_classes, n_channels)\n",
    "    if phase == 'train':\n",
    "        my_resnet.train()\n",
    "        torch_resnet.train()\n",
    "    elif phase == 'eval':\n",
    "        my_resnet.eval()\n",
    "        torch_resnet.eval()\n",
    "    else:\n",
    "        raise Exception(\"unknown phase\") \n",
    "\n",
    "    my_resnet.clone_weights_from_torch(torch_resnet)\n",
    "\n",
    "    my_out = my_resnet.bn1.forward(input_data)\n",
    "\n",
    "    torch_out = torch_resnet.bn1(input_data_torch)\n",
    "    torch_out_np = torch_out.detach().numpy()\n",
    "\n",
    "    torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "    torch_input_g = input_data_torch.grad.detach().numpy()\n",
    "\n",
    "    my_input_g = my_resnet.bn1.backward(output_gradient)\n",
    "\n",
    "    atol=1e-6\n",
    "\n",
    "    print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "    print(\"input gradients all close:\", np.allclose(my_input_g.flatten(), torch_input_g.flatten(), atol=atol))\n",
    "    print()\n",
    "\n",
    "print(\"It's ok that gradients don't match in eval mode. \"\\\n",
    "      \"Batch norm's backward is different in train and eval mode.\\n\"\\\n",
    "      \"I didn't implement eval backward since it won't ever be used in training the network\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a6997b3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "57fc2d31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "phase: train\n",
      "output all close: True\n",
      "input gradients all close: False\n",
      "\n",
      "phase: eval\n",
      "output all close: True\n",
      "input gradients all close: False\n",
      "\n",
      "It's ok that gradients don't match in eval mode. Batch norm's backward is different in train and eval mode.\n",
      "I didn't implement eval backward since it won't ever be used in training the network\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "conv2_x test\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "batch_size = 10\n",
    "height = width = 8\n",
    "n_channels = 64\n",
    "\n",
    "for phase in ['train', 'eval']:\n",
    "    print(f\"phase: {phase}\")\n",
    "\n",
    "    input_data = np.random.rand(batch_size, n_channels, height, width).astype(np.float32)\n",
    "    input_data_torch = torch.from_numpy(input_data).float()\n",
    "    input_data_torch.requires_grad = True\n",
    "    output_gradient = np.random.rand(batch_size, 256, 8, 8).astype(np.float32)\n",
    "\n",
    "    torch_resnet = resnet101_torch(10, 1)\n",
    "    my_resnet = resnet101(10, 1)\n",
    "\n",
    "    if phase == 'train':\n",
    "        my_resnet.train()\n",
    "        torch_resnet.train()\n",
    "    elif phase == 'eval':\n",
    "        my_resnet.eval()\n",
    "        torch_resnet.eval()\n",
    "    else:\n",
    "        raise Exception(\"unknown phase\") \n",
    "\n",
    "\n",
    "    my_resnet.clone_weights_from_torch(torch_resnet)\n",
    "\n",
    "    my_out = my_resnet.conv2_x.forward(input_data)\n",
    "\n",
    "    torch_out = torch_resnet.conv2_x(input_data_torch)\n",
    "    torch_out_np = torch_out.detach().numpy()\n",
    "\n",
    "    torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "    torch_input_g = input_data_torch.grad.detach().numpy()\n",
    "\n",
    "    my_input_g = my_resnet.conv2_x.backward(output_gradient)\n",
    "\n",
    "    atol=1e-3\n",
    "\n",
    "    print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "    print(\"input gradients all close:\", np.allclose(my_input_g, torch_input_g, atol=atol))\n",
    "    print()\n",
    "\n",
    "print(\"It's ok that gradients don't match in eval mode. \"\\\n",
    "      \"Batch norm's backward is different in train and eval mode.\\n\"\\\n",
    "      \"I didn't implement eval backward since it won't ever be used in training the network\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "9234786a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output all close: True\n",
      "input gradients all close: False\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "conv3_x test\n",
    "\"\"\"\n",
    "\n",
    "batch_size = 10\n",
    "height = width = 8\n",
    "n_channels = 256\n",
    "\n",
    "input_data = np.random.rand(batch_size, n_channels, height, width).astype(np.float32)\n",
    "input_data_torch = torch.from_numpy(input_data).float()\n",
    "input_data_torch.requires_grad = True\n",
    "output_gradient = np.random.rand(batch_size, 512, 4, 4).astype(np.float32)\n",
    "\n",
    "torch_resnet = resnet101_torch(10, 1)\n",
    "my_resnet = resnet101(10, 1)\n",
    "\n",
    "my_resnet.clone_weights_from_torch(torch_resnet)\n",
    "\n",
    "my_out = my_resnet.conv3_x.forward(input_data)\n",
    "\n",
    "torch_out = torch_resnet.conv3_x(input_data_torch)\n",
    "torch_out_np = torch_out.detach().numpy()\n",
    "\n",
    "torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "torch_input_g = input_data_torch.grad.detach().numpy()\n",
    "\n",
    "my_input_g = my_resnet.conv3_x.backward(output_gradient)\n",
    "\n",
    "atol=1e-3\n",
    "\n",
    "print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "print(\"input gradients all close:\", np.allclose(my_input_g, torch_input_g, atol=atol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "c31786ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output all close: True\n",
      "input gradients all close: False\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "conv4_x test\n",
    "\"\"\"\n",
    "\n",
    "batch_size = 10\n",
    "height = width = 4\n",
    "n_channels = 512\n",
    "\n",
    "input_data = np.random.rand(batch_size, n_channels, height, width).astype(np.float32)\n",
    "input_data_torch = torch.from_numpy(input_data).float()\n",
    "input_data_torch.requires_grad = True\n",
    "output_gradient = np.random.rand(batch_size, 1024, height//2, width//2).astype(np.float32)\n",
    "\n",
    "torch_resnet = resnet101_torch(10, 1)\n",
    "my_resnet = resnet101(10, 1)\n",
    "\n",
    "my_resnet.train()\n",
    "torch_resnet.train()\n",
    "\n",
    "my_resnet.clone_weights_from_torch(torch_resnet)\n",
    "\n",
    "my_out = my_resnet.conv4_x.forward(input_data)\n",
    "\n",
    "torch_out = torch_resnet.conv4_x(input_data_torch)\n",
    "\n",
    "torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "torch_input_g = input_data_torch.grad.detach().numpy()\n",
    "torch_out_np = torch_out.detach().numpy()\n",
    "\n",
    "\n",
    "my_input_g = my_resnet.conv4_x.backward(output_gradient)\n",
    "\n",
    "atol=1e-3\n",
    "\n",
    "print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "print(\"input gradients all close:\", np.allclose(my_input_g, torch_input_g, atol=atol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "e639853e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([-92.64273227,   0.        ,  81.38443665, ...,   0.        ,\n",
       "          0.        ,   0.        ]),\n",
       " array([-92.29285,   0.     ,  81.22645, ...,   0.     ,   0.     ,\n",
       "          0.     ], dtype=float32))"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_input_g.flatten(), torch_input_g.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "b704ef7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output all close: True\n",
      "input gradients all close: True\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "conv5_x test\n",
    "\"\"\"\n",
    "\n",
    "batch_size = 10\n",
    "height = width = 2\n",
    "n_channels = 1024\n",
    "\n",
    "input_data = np.random.rand(batch_size, n_channels, height, width).astype(np.float32)\n",
    "input_data_torch = torch.from_numpy(input_data).float()\n",
    "input_data_torch.requires_grad = True\n",
    "output_gradient = np.random.rand(batch_size, 2048, 1, 1).astype(np.float32)\n",
    "\n",
    "torch_resnet = resnet101_torch(10, 1)\n",
    "my_resnet = resnet101(10, 1)\n",
    "\n",
    "my_resnet.clone_weights_from_torch(torch_resnet)\n",
    "\n",
    "my_out = my_resnet.conv5_x.forward(input_data)\n",
    "\n",
    "torch_out = torch_resnet.conv5_x(input_data_torch)\n",
    "torch_out_np = torch_out.detach().numpy()\n",
    "\n",
    "torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "torch_input_g = input_data_torch.grad.detach().numpy()\n",
    "\n",
    "my_input_g = my_resnet.conv5_x.backward(output_gradient)\n",
    "\n",
    "atol=1e-3\n",
    "\n",
    "print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "print(\"input gradients all close:\", np.allclose(my_input_g, torch_input_g, atol=atol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "affa6502",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output all close: True\n",
      "input gradients all close: True\n",
      "before adam weights gradients all close: True\n",
      "before adam bias gradients all close: True\n",
      "[[1.3891779  1.7336051  1.2286544 ]\n",
      " [1.3048313  0.99781084 1.0630406 ]\n",
      " [1.6351756  1.4931997  1.3154104 ]\n",
      " [1.5244397  1.4477336  1.157729  ]\n",
      " [0.9432854  0.58816004 0.951926  ]\n",
      " [1.6254252  2.1075394  1.3502117 ]] \n",
      " [[1.3891779  1.733605   1.2286544 ]\n",
      " [1.3048313  0.99781084 1.0630406 ]\n",
      " [1.6351756  1.4931998  1.3154104 ]\n",
      " [1.5244397  1.4477336  1.1577291 ]\n",
      " [0.9432854  0.58816004 0.95192605]\n",
      " [1.6254252  2.1075392  1.3502117 ]]\n",
      "after adam weights gradients all close: True\n",
      "after adam bias gradients all close: True\n",
      "[[1.3891779  1.7336051  1.2286544 ]\n",
      " [1.3048313  0.99781084 1.0630406 ]\n",
      " [1.6351756  1.4931997  1.3154104 ]\n",
      " [1.5244397  1.4477336  1.157729  ]\n",
      " [0.9432854  0.58816004 0.951926  ]\n",
      " [1.6254252  2.1075394  1.3502117 ]] \n",
      " [[1.3891779  1.733605   1.2286544 ]\n",
      " [1.3048313  0.99781084 1.0630406 ]\n",
      " [1.6351756  1.4931998  1.3154104 ]\n",
      " [1.5244397  1.4477336  1.1577291 ]\n",
      " [0.9432854  0.58816004 0.95192605]\n",
      " [1.6254252  2.1075392  1.3502117 ]]\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "AdamOptimizer test\n",
    "\"\"\"\n",
    "\n",
    "from torch.optim import Adam as Adam_torch\n",
    "\n",
    "n_input_features = 6\n",
    "n_output_features = 3\n",
    "batch_size = 5\n",
    "input_data = np.random.rand(batch_size, n_input_features).astype(np.float32)\n",
    "input_data_torch = torch.from_numpy(input_data).float()\n",
    "input_data_torch.requires_grad = True\n",
    "output_gradient = np.random.rand(batch_size, n_output_features).astype(np.float32)\n",
    "\n",
    "torch_fc = torch.nn.Linear(n_input_features, n_output_features)\n",
    "torch_out = torch_fc(input_data_torch)\n",
    "torch_out_np = torch_out.detach().numpy()\n",
    "torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "torch_wg = torch_fc.weight.grad.detach().numpy().T\n",
    "torch_bg = torch_fc.bias.grad.detach().numpy().reshape(-1, 1).T\n",
    "torch_input_g = input_data_torch.grad.detach().numpy()\n",
    "\n",
    "\n",
    "my_fc = FullyConnectedLayer(n_input_features, n_output_features)\n",
    "my_fc.weights = torch_fc.weight.detach().numpy().T\n",
    "my_fc.bias = torch_fc.bias.detach().numpy().reshape(-1, 1).T\n",
    "my_out = my_fc.forward(input_data)\n",
    "my_input_g = my_fc.backward(output_gradient)\n",
    "my_wg = my_fc.weights_gradient\n",
    "my_bg = my_fc.bias_gradient\n",
    "\n",
    "atol=1e-3\n",
    "\n",
    "print(\"output all close:\", np.allclose(my_out, torch_out_np, atol=atol))\n",
    "print(\"input gradients all close:\", np.allclose(my_input_g, torch_input_g, atol=atol))\n",
    "print(\"before adam weights gradients all close:\", np.allclose(my_wg, torch_wg, atol=atol))\n",
    "print(\"before adam bias gradients all close:\", np.allclose(my_bg, torch_bg, atol=atol))\n",
    "print(my_wg, \"\\n\", torch_wg)\n",
    "\n",
    "my_adam = AdamOptimizer(my_fc.get_trainable_layers(), 0.001, 0.9, 0.999, 1e-8)\n",
    "my_adam.step()\n",
    "\n",
    "torch_adam = Adam_torch(torch_fc.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-8)\n",
    "torch_adam.step()\n",
    "\n",
    "print(\"after adam weights gradients all close:\", np.allclose(my_fc.weights_gradient, torch_fc.weight.grad.detach().numpy().T, atol=atol))\n",
    "print(\"after adam bias gradients all close:\", np.allclose(my_fc.bias_gradient, torch_fc.bias.grad.detach().numpy().reshape(-1, 1).T, atol=atol))\n",
    "print(my_fc.weights_gradient, \"\\n\", torch_fc.weight.grad.detach().numpy().T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "5bca9fe7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batch_size: 1\n",
      "my_conv forward time: 0.5448141098022461\n",
      "my_conv_with_loops forward time: 3.00417423248291\n",
      "torch_conv forward time: 0.07560348510742188\n",
      "my_conv backward time: 0.684333324432373\n",
      "my_conv_with_loops backward time: 2.6092746257781982\n",
      "torch_conv backward time: 0.17493200302124023\n",
      "\n",
      "batch_size: 2\n",
      "my_conv forward time: 0.5803859233856201\n",
      "my_conv_with_loops forward time: 2.1322405338287354\n",
      "torch_conv forward time: 0.15914392471313477\n",
      "my_conv backward time: 1.183729648590088\n",
      "my_conv_with_loops backward time: 4.955396413803101\n",
      "torch_conv backward time: 0.365201473236084\n",
      "\n",
      "batch_size: 4\n",
      "my_conv forward time: 1.063096284866333\n",
      "my_conv_with_loops forward time: 2.2665672302246094\n",
      "torch_conv forward time: 0.16541314125061035\n",
      "my_conv backward time: 2.5409464836120605\n",
      "my_conv_with_loops backward time: 9.793752670288086\n",
      "torch_conv backward time: 0.414461612701416\n",
      "\n",
      "batch_size: 8\n",
      "my_conv forward time: 1.992255687713623\n",
      "my_conv_with_loops forward time: 2.455413341522217\n",
      "torch_conv forward time: 0.20453190803527832\n",
      "my_conv backward time: 4.4292519092559814\n",
      "my_conv_with_loops backward time: 21.643410205841064\n",
      "torch_conv backward time: 0.5241467952728271\n",
      "\n",
      "batch_size: 16\n",
      "my_conv forward time: 4.281964540481567\n",
      "my_conv_with_loops forward time: 3.1336183547973633\n",
      "torch_conv forward time: 0.3031580448150635\n",
      "my_conv backward time: 9.075878858566284\n",
      "my_conv_with_loops backward time: 41.86888933181763\n",
      "torch_conv backward time: 0.500030517578125\n",
      "\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Conv2d vs Conv2dWithLoops vs torch.nn.Conv2d time comparison forward and backward\n",
    "\"\"\"\n",
    "\n",
    "import time\n",
    "\n",
    "n_input_channels = 4\n",
    "n_output_channels = 2\n",
    "width = 3\n",
    "height = 5\n",
    "\n",
    "kernel_size = 3\n",
    "stride = 1\n",
    "padding = 3\n",
    "\n",
    "output_width = (width + 2 * padding - kernel_size) // stride + 1\n",
    "output_height = (height + 2 * padding - kernel_size) // stride + 1\n",
    "\n",
    "for batch_size in [1, 2, 4, 8, 16]:\n",
    "    print(\"batch_size:\", batch_size)\n",
    "\n",
    "    input_data = np.random.rand(batch_size, n_input_channels, height, width).astype(np.float32)\n",
    "    input_data_torch = torch.from_numpy(input_data).float()\n",
    "    input_data_torch.requires_grad = True\n",
    "    output_gradient = np.random.rand(batch_size, n_output_channels, output_height, output_width).astype(np.float32)\n",
    "\n",
    "    torch_conv = torch.nn.Conv2d(n_input_channels, n_output_channels, kernel_size, stride, padding)\n",
    "\n",
    "    my_conv_with_loops = Conv2dWithLoops(n_input_channels, n_output_channels, kernel_size, stride, padding)\n",
    "    my_conv_with_loops.weights = torch_conv.weight.detach().numpy()\n",
    "\n",
    "    my_conv = Conv2d(n_input_channels, n_output_channels, kernel_size, stride, padding)\n",
    "    my_conv.weights = torch_conv.weight.detach().numpy()\n",
    "\n",
    "    n_iterations = 1000\n",
    "\n",
    "    start = time.time()\n",
    "    for i in range(n_iterations):\n",
    "        my_out = my_conv.forward(input_data)\n",
    "    end = time.time()\n",
    "    print(f\"my_conv forward time: {end - start}\")\n",
    "\n",
    "    start = time.time()\n",
    "    for i in range(n_iterations):\n",
    "        my_out_with_loops = my_conv_with_loops.forward(input_data)\n",
    "    end = time.time()\n",
    "    print(f\"my_conv_with_loops forward time: {end - start}\")\n",
    "\n",
    "    start = time.time()\n",
    "    for i in range(n_iterations):\n",
    "        torch_out = torch_conv(input_data_torch)\n",
    "    end = time.time()\n",
    "    print(f\"torch_conv forward time: {end - start}\")\n",
    "\n",
    "    start = time.time()\n",
    "    for i in range(n_iterations):\n",
    "        my_input_g = my_conv.backward_as_matrix_multiplication(output_gradient)\n",
    "    end = time.time()\n",
    "    print(f\"my_conv backward time: {end - start}\")\n",
    "\n",
    "    start = time.time()\n",
    "    for i in range(n_iterations):\n",
    "        my_input_g_with_loops = my_conv_with_loops.backward(output_gradient)\n",
    "    end = time.time()\n",
    "    print(f\"my_conv_with_loops backward time: {end - start}\")\n",
    "\n",
    "    start = time.time()\n",
    "    for i in range(n_iterations):\n",
    "        torch_out.backward(torch.tensor(output_gradient), retain_graph=True)\n",
    "    end = time.time()\n",
    "    print(f\"torch_conv backward time: {end - start}\")\n",
    "\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e07955d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e89c2b9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "673fcb0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GlobalAveragePooling2D:\n",
    "    def __init__(self):\n",
    "        # self.input_shape = None\n",
    "        pass\n",
    "\n",
    "    def forward(self, input_):\n",
    "        self.input_ = input_\n",
    "        self.input_shape = input_.shape\n",
    "        return np.mean(input_, axis=(2, 3))\n",
    "    \n",
    "    def backward(self, d_J_d_out):\n",
    "        d_out_d_in = np.ones(self.input_shape) / np.prod(self.input_shape[2:])\n",
    "        d_J_d_out = d_J_d_out[:, :, np.newaxis, np.newaxis]\n",
    "        return d_out_d_in * d_J_d_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "6b992530",
   "metadata": {},
   "outputs": [],
   "source": [
    "gap = GlobalAveragePooling2D()\n",
    "\n",
    "torch_gap = torch.nn.Sequential(\n",
    "    torch.nn.AdaptiveAvgPool2d((1, 1)),\n",
    "    torch.nn.Flatten(start_dim=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "7e42574d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make torch_gap require grad\n",
    "torch_gap[0].requires_grad = True\n",
    "torch_gap[1].requires_grad = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "c80b6c47",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (1): Flatten(start_dim=1, end_dim=-1)\n",
       ")"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch_gap.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "19953027",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "out all close: True\n",
      "d_J_d_in all close: True\n"
     ]
    }
   ],
   "source": [
    "input_ = np.random.rand(2, 3, 4, 5)\n",
    "\n",
    "input_torch = torch.from_numpy(input_).float()\n",
    "\n",
    "input_torch.requires_grad = True\n",
    "\n",
    "d_J_d_out = np.random.rand(2, 3)\n",
    "\n",
    "out = gap.forward(input_)\n",
    "\n",
    "torch_out = torch_gap(input_torch)\n",
    "\n",
    "d_J_d_in = gap.backward(d_J_d_out)\n",
    "\n",
    "torch_out.backward(torch.tensor(d_J_d_out), retain_graph=True)\n",
    "\n",
    "print(\"out all close:\", np.allclose(out, torch_out.detach().numpy()))\n",
    "print(\"d_J_d_in all close:\", np.allclose(d_J_d_in, input_torch.grad.detach().numpy()))"
   ]
  }
 ],
 "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.10.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "2bbf79ea567ebf64b92b9cf68f2b08a0cf8db5ff8a7f29f56d99d802810464d9"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}