How to freeze parameters in a model using pytorch

freeze_parameters.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "class PreModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(PreModel, self).__init__()\n",
    "        self.c1 = nn.Conv1d(1, 2, 3, bias=False)\n",
    "    def forward(self, x):\n",
    "        return self.c1(x)\n",
    "class PostModel(nn.Module):\n",
    "    def __init__(self, train_eye_net=True):\n",
    "        super(PostModel, self).__init__()\n",
    "        self.pre = PreModel()\n",
    "        for k in self.pre.parameters():\n",
    "            k.requires_grad = train_eye_net\n",
    "        self.c2 = nn.Conv1d(2, 3, 4, bias=False)\n",
    "    def forward(self, x):\n",
    "        eye_out = self.pre(x)\n",
    "        return self.c2(eye_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "post = PostModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_param(in_model):\n",
    "    return {\n",
    "        tuple(k.size()): k.detach().numpy().copy()\n",
    "        for k in in_model.parameters()\n",
    "    }\n",
    "old_param = get_param(post)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = torch.nn.MSELoss(reduction='sum')\n",
    "optimizer = torch.optim.SGD(post.parameters(), lr=1e-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Normal training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(2, 1, 10)\n",
    "y_pred = post(x)\n",
    "y_true = torch.zeros_like(y_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(10):\n",
    "    x = torch.randn(2, 1, 10)\n",
    "    y_pred = post(x)\n",
    "    loss = criterion(y_pred, y_true)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_match(in_model, in_param):\n",
    "    for k in in_model.parameters():\n",
    "        dim = tuple(k.size())\n",
    "        new_val = k.detach().numpy()\n",
    "        old_val = in_param[dim]\n",
    "        if np.allclose(new_val, old_val):\n",
    "            print(dim, 'matches')\n",
    "        else:\n",
    "            print(dim, 'doesnt match', np.mean(np.abs(new_val-old_val)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 1, 3) doesnt match 0.20992236\n",
      "(3, 2, 4) doesnt match 0.101107724\n"
     ]
    }
   ],
   "source": [
    "check_match(post, old_param)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Disable require_grad on the first module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "post_fix = PostModel(train_eye_net=False)\n",
    "old_param = get_param(post_fix)\n",
    "optimizer = torch.optim.SGD(post_fix.parameters(), lr=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(10):\n",
    "    x = torch.randn(2, 1, 10)\n",
    "    y_pred = post_fix(x)\n",
    "    loss = criterion(y_pred, y_true)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 1, 3) matches\n",
      "(3, 2, 4) doesnt match 29.820532\n"
     ]
    }
   ],
   "source": [
    "check_match(post_fix, old_param)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}