Graphics generated by torchviz

CherkNevis.ipynb

@@ -237,9 +237,67 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "id": "23476fd8-6145-4f7d-81f7-adbc4855e452",
    "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'LSTM_model.svg'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from torchviz import make_dot\n",
+    "\n",
+    "# Define the LSTM model\n",
+    "class LSTMNet(nn.Module):\n",
+    "    def __init__(self, input_size, hidden_size, output_size, n_layers=1):\n",
+    "        super(LSTMNet, self).__init__()\n",
+    "        self.hidden_size = hidden_size\n",
+    "        self.n_layers = n_layers\n",
+    "        self.lstm = nn.LSTM(input_size, hidden_size, n_layers, batch_first=True)\n",
+    "        self.fc = nn.Linear(hidden_size, output_size)\n",
+    "    \n",
+    "    def forward(self, x):\n",
+    "        h0 = torch.zeros(self.n_layers, x.size(0), self.hidden_size).to(x.device)\n",
+    "        c0 = torch.zeros(self.n_layers, x.size(0), self.hidden_size).to(x.device)\n",
+    "        out, _ = self.lstm(x, (h0, c0))\n",
+    "        out = self.fc(out[:, -1, :])\n",
+    "        return out\n",
+    "\n",
+    "# Hyperparameters\n",
+    "input_size = 1\n",
+    "hidden_size = 50\n",
+    "output_size = 1\n",
+    "num_layers = 1\n",
+    "\n",
+    "# Initialize the model\n",
+    "model = LSTMNet(input_size, hidden_size, output_size, num_layers)\n",
+    "\n",
+    "# Create a dummy input\n",
+    "dummy_input = torch.randn(1, 10, input_size)  # Batch size = 1, Sequence length = 10, Input size = 1\n",
+    "\n",
+    "# Generate the computation graph\n",
+    "output = model(dummy_input)\n",
+    "dot = make_dot(output, params=dict(model.named_parameters()))\n",
+    "dot.format = 'svg'\n",
+    "dot.render(\"LSTM_model\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "18b89e50-14a6-4bd1-9841-e876ae2bc983",
+   "metadata": {},
    "outputs": [],
    "source": []
   }

LSTM_model

@@ -0,0 +1,50 @@
+digraph {
+	graph [size="12,12"]
+	node [align=left fontname=monospace fontsize=10 height=0.2 ranksep=0.1 shape=box style=filled]
+	2283802772080 [label="
+ (1, 1)" fillcolor=darkolivegreen1]
+	2283802592880 [label=AddmmBackward0]
+	2283802592688 -> 2283802592880
+	2283802772000 [label="fc.bias
+ (1)" fillcolor=lightblue]
+	2283802772000 -> 2283802592688
+	2283802592688 [label=AccumulateGrad]
+	2283802595280 -> 2283802592880
+	2283802595280 [label=SliceBackward0]
+	2283802593312 -> 2283802595280
+	2283802593312 [label=SelectBackward0]
+	2283802592112 -> 2283802593312
+	2283802592112 [label=SliceBackward0]
+	2283802593504 -> 2283802592112
+	2283802593504 [label=TransposeBackward0]
+	2283802592928 -> 2283802593504
+	2283802592928 [label=MkldnnRnnLayerBackward0]
+	2283802593408 -> 2283802592928
+	2283656326896 [label="lstm.weight_ih_l0
+ (200, 1)" fillcolor=lightblue]
+	2283656326896 -> 2283802593408
+	2283802593408 [label=AccumulateGrad]
+	2283802593744 -> 2283802592928
+	2283656326816 [label="lstm.weight_hh_l0
+ (200, 50)" fillcolor=lightblue]
+	2283656326816 -> 2283802593744
+	2283802593744 [label=AccumulateGrad]
+	2283802593648 -> 2283802592928
+	2283656326976 [label="lstm.bias_ih_l0
+ (200)" fillcolor=lightblue]
+	2283656326976 -> 2283802593648
+	2283802593648 [label=AccumulateGrad]
+	2283802594032 -> 2283802592928
+	2283656327056 [label="lstm.bias_hh_l0
+ (200)" fillcolor=lightblue]
+	2283656327056 -> 2283802594032
+	2283802594032 [label=AccumulateGrad]
+	2283802594992 -> 2283802592880
+	2283802594992 [label=TBackward0]
+	2283802593024 -> 2283802594992
+	2283802771920 [label="fc.weight
+ (1, 50)" fillcolor=lightblue]
+	2283802771920 -> 2283802593024
+	2283802593024 [label=AccumulateGrad]
+	2283802592880 -> 2283802772080
+}