{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e3a42e6e-0232-4c4a-b88b-7805d7a41347",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "61fbb71c-1aa8-4fb7-9e1f-ee2023d2517b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        # 1 input image channel, 6 output channels, 5x5 square convolution\n",
    "        # kernel\n",
    "        self.conv1 = nn.Conv2d(1, 6, 5)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        # an affine operation: y = Wx + b\n",
    "        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5*5 from image dimension\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.fc3 = nn.Linear(84, 10)\n",
    "\n",
    "    def forward(self, input):\n",
    "        # Convolution layer C1: 1 input image channel, 6 output channels,\n",
    "        # 5x5 square convolution, it uses RELU activation function, and\n",
    "        # outputs a Tensor with size (N, 6, 28, 28), where N is the size of the batch\n",
    "        c1 = F.relu(self.conv1(input))\n",
    "        # Subsampling layer S2: 2x2 grid, purely functional,\n",
    "        # this layer does not have any parameter, and outputs a (N, 6, 14, 14) Tensor\n",
    "        s2 = F.max_pool2d(c1, (2, 2))\n",
    "        # Convolution layer C3: 6 input channels, 16 output channels,\n",
    "        # 5x5 square convolution, it uses RELU activation function, and\n",
    "        # outputs a (N, 16, 10, 10) Tensor\n",
    "        c3 = F.relu(self.conv2(s2))\n",
    "        # Subsampling layer S4: 2x2 grid, purely functional,\n",
    "        # this layer does not have any parameter, and outputs a (N, 16, 5, 5) Tensor\n",
    "        s4 = F.max_pool2d(c3, 2)\n",
    "        # Flatten operation: purely functional, outputs a (N, 400) Tensor\n",
    "        s4 = torch.flatten(s4, 1)\n",
    "        # Fully connected layer F5: (N, 400) Tensor input,\n",
    "        # and outputs a (N, 120) Tensor, it uses RELU activation function\n",
    "        f5 = F.relu(self.fc1(s4))\n",
    "        # Fully connected layer F6: (N, 120) Tensor input,\n",
    "        # and outputs a (N, 84) Tensor, it uses RELU activation function\n",
    "        f6 = F.relu(self.fc2(f5))\n",
    "        # Fully connected layer OUTPUT: (N, 84) Tensor input, and\n",
    "        # outputs a (N, 10) Tensor\n",
    "        output = self.fc3(f6)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f798d5f9-7013-44a3-9672-c740050c4504",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Net(\n",
      "  (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))\n",
      "  (fc1): Linear(in_features=400, out_features=120, bias=True)\n",
      "  (fc2): Linear(in_features=120, out_features=84, bias=True)\n",
      "  (fc3): Linear(in_features=84, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "net = Net()\n",
    "print(net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "045d621d-25cd-4a6f-a211-952cc624a312",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10\n",
      "torch.Size([6, 1, 5, 5])\n"
     ]
    }
   ],
   "source": [
    "params = list(net.parameters())\n",
    "print(len(params))\n",
    "print(params[0].size())  # conv1's .weight\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0b48593d-0eed-4c33-bd26-0fc768fc72dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.0708, -0.0528, -0.1110,  0.0017,  0.0315, -0.0709, -0.0408, -0.1356,\n",
      "         -0.1438,  0.0293]], grad_fn=<AddmmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "input = torch.randn(1, 1, 32, 32)\n",
    "out = net(input)\n",
    "print(out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e150e04b-346a-41b0-8688-1a951d7b5272",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}