{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Extract layer saturation\nExtract layer saturation with Delve.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\nfrom tqdm import trange\n\nfrom delve import CheckLayerSat\n\n\nclass TwoLayerNet(torch.nn.Module):\n    def __init__(self, D_in, H, D_out):\n        super(TwoLayerNet, self).__init__()\n        self.linear1 = torch.nn.Linear(D_in, H)\n        self.linear2 = torch.nn.Linear(H, D_out)\n\n    def forward(self, x):\n        h_relu = self.linear1(x).clamp(min=0)\n        y_pred = self.linear2(h_relu)\n        return y_pred\n\n\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\ntorch.manual_seed(1)\n\nfor h in [3, 32]:\n    # N is batch size; D_in is input dimension;\n    # H is hidden dimension; D_out is output dimension.\n    N, D_in, H, D_out = 64, 1000, h, 10\n\n    # Create random Tensors to hold inputs and outputs\n    x = torch.randn(N, D_in)\n    y = torch.randn(N, D_out)\n    x_test = torch.randn(N, D_in)\n    y_test = torch.randn(N, D_out)\n\n    # You can watch specific layers by handing them to delve as a list.\n    # Also, you can hand over the entire Module-object to delve and let delve search for recordable layers.\n    model = TwoLayerNet(D_in, H, D_out)\n\n    x, y, model = x.to(device), y.to(device), model.to(device)\n    x_test, y_test = x_test.to(device), y_test.to(device)\n\n    layers = [model.linear1, model.linear2]\n    stats = CheckLayerSat('regression/h{}'.format(h),\n                          save_to=\"plotcsv\",\n                          modules=layers,\n                          device=device,\n                          stats=[\"lsat\", \"lsat_eval\"])\n\n    loss_fn = torch.nn.MSELoss(reduction='sum')\n    optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)\n    steps_iter = trange(2000, desc='steps', leave=True, position=0)\n    steps_iter.write(\"{:^80}\".format(\n        \"Regression - TwoLayerNet - Hidden layer size {}\".format(h)))\n    for step in steps_iter:\n        # training step\n        model.train()\n        y_pred = model(x)\n        loss = loss_fn(y_pred, y)\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        # test step\n        model.eval()\n        y_pred = model(x_test)\n        loss_test = loss_fn(y_pred, y_test)\n\n        # update statistics\n        steps_iter.set_description('loss=%g' % loss.item())\n        stats.add_scalar(\"train-loss\", loss.item())\n        stats.add_scalar(\"test-loss\", loss_test.item())\n\n        stats.add_saturations()\n    steps_iter.write('\\n')\n    stats.close()\n    steps_iter.close()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}