tutorials/2_Visualization.ipynb

{
  "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization in TorchOpt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[<img align=\"left\" src=\"https://colab.research.google.com/assets/colab-badge.svg\">](https://colab.research.google.com/github/metaopt/torchopt/blob/main/tutorials/2_Visualization.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In [PyTorch](https://pytorch.org), if the attribute `requires_grad` of a tensor is `True`, the computation graph will be created if we use the tensor to do any operations. The computation graph is implemented like link-list -- `Tensor`s are nodes and they are linked by their attribute `gran_fn`. [PyTorchViz](https://github.com/szagoruyko/pytorchviz) is a Python package that uses [Graphviz](https://graphviz.org) as a backend for plotting computation graphs. TorchOpt use PyTorchViz as the blueprint and provide more easy-to-use visualization functions on the premise of supporting all its functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start with a simple multiplication computation graph. We declared the variable `x` with flag `requires_grad=True` and compute `y = 2 * x`. Then we visualize the computation graph of `y`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<graphviz.graphs.Digraph object at 0x7fd0a30377f0>\n"
     ]
    },
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Generated by graphviz version 2.42.3 (20191010.1750)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"109pt\" height=\"214pt\"\n viewBox=\"0.00 0.00 109.00 214.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 210)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-210 105,-210 105,4 -4,4\"/>\n<!-- 140534064715952 -->\n<g id=\"node1\" class=\"node\">\n<title>140534064715952</title>\n<polygon fill=\"#caff70\" stroke=\"black\" points=\"77.5,-30 23.5,-30 23.5,0 77.5,0 77.5,-30\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-18\" font-family=\"monospace\" font-size=\"10.00\">y</text>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-7\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140534064838304 -->\n<g id=\"node2\" class=\"node\">\n<title>140534064838304</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"95,-85 6,-85 6,-66 95,-66 95,-85\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-73\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140534064838304&#45;&gt;140534064715952 -->\n<g id=\"edge3\" class=\"edge\">\n<title>140534064838304&#45;&gt;140534064715952</title>\n<path fill=\"none\" stroke=\"black\" d=\"M50.5,-65.87C50.5,-59.11 50.5,-49.35 50.5,-40.26\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54,-40.11 50.5,-30.11 47,-40.11 54,-40.11\"/>\n</g>\n<!-- 140534064837776 -->\n<g id=\"node3\" class=\"node\">\n<title>140534064837776</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"101,-140 0,-140 0,-121 101,-121 101,-140\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-128\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140534064837776&#45;&gt;140534064838304 -->\n<g id=\"edge1\" class=\"edge\">\n<title>140534064837776&#45;&gt;140534064838304</title>\n<path fill=\"none\" stroke=\"black\" d=\"M50.5,-120.75C50.5,-113.8 50.5,-103.85 50.5,-95.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54,-95.09 50.5,-85.09 47,-95.09 54,-95.09\"/>\n</g>\n<!-- 140534064714832 -->\n<g id=\"node4\" class=\"node\">\n<title>140534064714832</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"77.5,-206 23.5,-206 23.5,-176 77.5,-176 77.5,-206\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-194\" font-family=\"monospace\" font-size=\"10.00\">x</text>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-183\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140534064714832&#45;&gt;140534064837776 -->\n<g id=\"edge2\" class=\"edge\">\n<title>140534064714832&#45;&gt;140534064837776</title>\n<path fill=\"none\" stroke=\"black\" d=\"M50.5,-175.84C50.5,-168.21 50.5,-158.7 50.5,-150.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54,-150.27 50.5,-140.27 47,-150.27 54,-150.27\"/>\n</g>\n</g>\n</svg>\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import torchopt\n",
    "\n",
    "\n",
    "x = torch.tensor(1.0, requires_grad=True)\n",
    "y = 2 * x\n",
    "display(torchopt.visual.make_dot(y, params={'x': x, 'y': y}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The figure shows `y` is connected by the multiplication edge. The gradient of `y` will flow through the multiplication backward function then accumulated on `x`. Note that we pass a dictionary for adding node labels."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then let's plot a neural network. Note that we can pass the generator returned by method `named_parameters` for adding node labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<graphviz.graphs.Digraph object at 0x7fd00fd56e20>\n"
     ]
    },
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Generated by graphviz version 2.42.3 (20191010.1750)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"216pt\" height=\"335pt\"\n viewBox=\"0.00 0.00 216.00 335.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 331)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-331 212,-331 212,4 -4,4\"/>\n<!-- 140534659780336 -->\n<g id=\"node1\" class=\"node\">\n<title>140534659780336</title>\n<polygon fill=\"#caff70\" stroke=\"black\" points=\"130.5,-30 76.5,-30 76.5,0 130.5,0 130.5,-30\"/>\n<text text-anchor=\"middle\" x=\"103.5\" y=\"-18\" font-family=\"monospace\" font-size=\"10.00\">loss</text>\n<text text-anchor=\"middle\" x=\"103.5\" y=\"-7\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140531595570768 -->\n<g id=\"node2\" class=\"node\">\n<title>140531595570768</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"160,-85 47,-85 47,-66 160,-66 160,-85\"/>\n<text text-anchor=\"middle\" x=\"103.5\" y=\"-73\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackward0</text>\n</g>\n<!-- 140531595570768&#45;&gt;140534659780336 -->\n<g id=\"edge7\" class=\"edge\">\n<title>140531595570768&#45;&gt;140534659780336</title>\n<path fill=\"none\" stroke=\"black\" d=\"M103.5,-65.87C103.5,-59.11 103.5,-49.35 103.5,-40.26\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"107,-40.11 103.5,-30.11 100,-40.11 107,-40.11\"/>\n</g>\n<!-- 140531595570576 -->\n<g id=\"node3\" class=\"node\">\n<title>140531595570576</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"154,-140 53,-140 53,-121 154,-121 154,-140\"/>\n<text text-anchor=\"middle\" x=\"103.5\" y=\"-128\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140531595570576&#45;&gt;140531595570768 -->\n<g id=\"edge1\" class=\"edge\">\n<title>140531595570576&#45;&gt;140531595570768</title>\n<path fill=\"none\" stroke=\"black\" d=\"M103.5,-120.75C103.5,-113.8 103.5,-103.85 103.5,-95.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"107,-95.09 103.5,-85.09 100,-95.09 107,-95.09\"/>\n</g>\n<!-- 140531595570528 -->\n<g id=\"node4\" class=\"node\">\n<title>140531595570528</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"101,-195 0,-195 0,-176 101,-176 101,-195\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-183\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140531595570528&#45;&gt;140531595570576 -->\n<g id=\"edge2\" class=\"edge\">\n<title>140531595570528&#45;&gt;140531595570576</title>\n<path fill=\"none\" stroke=\"black\" d=\"M59.25,-175.75C66.97,-168.03 78.4,-156.6 87.72,-147.28\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"90.31,-149.64 94.91,-140.09 85.36,-144.69 90.31,-149.64\"/>\n</g>\n<!-- 140531595583632 -->\n<g id=\"node5\" class=\"node\">\n<title>140531595583632</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"80,-261 21,-261 21,-231 80,-231 80,-261\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-249\" font-family=\"monospace\" font-size=\"10.00\">fc.bias</text>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-238\" font-family=\"monospace\" font-size=\"10.00\">(1)</text>\n</g>\n<!-- 140531595583632&#45;&gt;140531595570528 -->\n<g id=\"edge3\" class=\"edge\">\n<title>140531595583632&#45;&gt;140531595570528</title>\n<path fill=\"none\" stroke=\"black\" d=\"M50.5,-230.84C50.5,-223.21 50.5,-213.7 50.5,-205.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54,-205.27 50.5,-195.27 47,-205.27 54,-205.27\"/>\n</g>\n<!-- 140531595571104 -->\n<g id=\"node6\" class=\"node\">\n<title>140531595571104</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"196,-195 119,-195 119,-176 196,-176 196,-195\"/>\n<text text-anchor=\"middle\" x=\"157.5\" y=\"-183\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140531595571104&#45;&gt;140531595570576 -->\n<g id=\"edge4\" class=\"edge\">\n<title>140531595571104&#45;&gt;140531595570576</title>\n<path fill=\"none\" stroke=\"black\" d=\"M148.58,-175.75C140.72,-168.03 129.07,-156.6 119.58,-147.28\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"121.84,-144.6 112.25,-140.09 116.94,-149.59 121.84,-144.6\"/>\n</g>\n<!-- 140531595570432 -->\n<g id=\"node7\" class=\"node\">\n<title>140531595570432</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"208,-255.5 107,-255.5 107,-236.5 208,-236.5 208,-255.5\"/>\n<text text-anchor=\"middle\" x=\"157.5\" y=\"-243.5\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140531595570432&#45;&gt;140531595571104 -->\n<g id=\"edge5\" class=\"edge\">\n<title>140531595570432&#45;&gt;140531595571104</title>\n<path fill=\"none\" stroke=\"black\" d=\"M157.5,-236.37C157.5,-228.25 157.5,-215.81 157.5,-205.39\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"161,-205.17 157.5,-195.17 154,-205.17 161,-205.17\"/>\n</g>\n<!-- 140531595582816 -->\n<g id=\"node8\" class=\"node\">\n<title>140531595582816</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"193,-327 122,-327 122,-297 193,-297 193,-327\"/>\n<text text-anchor=\"middle\" x=\"157.5\" y=\"-315\" font-family=\"monospace\" font-size=\"10.00\">fc.weight</text>\n<text text-anchor=\"middle\" x=\"157.5\" y=\"-304\" font-family=\"monospace\" font-size=\"10.00\">(1, 5)</text>\n</g>\n<!-- 140531595582816&#45;&gt;140531595570432 -->\n<g id=\"edge6\" class=\"edge\">\n<title>140531595582816&#45;&gt;140531595570432</title>\n<path fill=\"none\" stroke=\"black\" d=\"M157.5,-296.8C157.5,-287.7 157.5,-275.79 157.5,-265.9\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"161,-265.84 157.5,-255.84 154,-265.84 161,-265.84\"/>\n</g>\n</g>\n</svg>\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Linear(dim, 1, bias=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc(x)\n",
    "\n",
    "\n",
    "dim = 5\n",
    "batch_size = 2\n",
    "net = Net(dim)\n",
    "xs = torch.ones((batch_size, dim))\n",
    "ys = torch.ones((batch_size, 1))\n",
    "pred = net(xs)\n",
    "loss = F.mse_loss(pred, ys)\n",
    "\n",
    "display(torchopt.visual.make_dot(loss, params=(net.named_parameters(), {'loss': loss})))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The computation graph of meta-learning algorithms will be much more complex. Our visualization tool allows users take as input the extracted network state for better visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<graphviz.graphs.Digraph object at 0x7fd00fd56640>\n"
     ]
    },
    {
     "data": {
      "image/svg+xml": "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n<!-- Generated by graphviz version 2.42.3 (20191010.1750)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"402pt\" height=\"995pt\"\n viewBox=\"0.00 0.00 402.00 995.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 991)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-991 398,-991 398,4 -4,4\"/>\n<!-- 140531595614064 -->\n<g id=\"node1\" class=\"node\">\n<title>140531595614064</title>\n<polygon fill=\"#caff70\" stroke=\"black\" points=\"298.5,-30 244.5,-30 244.5,0 298.5,0 298.5,-30\"/>\n<text text-anchor=\"middle\" x=\"271.5\" y=\"-18\" font-family=\"monospace\" font-size=\"10.00\">loss</text>\n<text text-anchor=\"middle\" x=\"271.5\" y=\"-7\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140531595567168 -->\n<g id=\"node2\" class=\"node\">\n<title>140531595567168</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"328,-85 215,-85 215,-66 328,-66 328,-85\"/>\n<text text-anchor=\"middle\" x=\"271.5\" y=\"-73\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackward0</text>\n</g>\n<!-- 140531595567168&#45;&gt;140531595614064 -->\n<g id=\"edge28\" class=\"edge\">\n<title>140531595567168&#45;&gt;140531595614064</title>\n<path fill=\"none\" stroke=\"black\" d=\"M271.5,-65.87C271.5,-59.11 271.5,-49.35 271.5,-40.26\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"275,-40.11 271.5,-30.11 268,-40.11 275,-40.11\"/>\n</g>\n<!-- 140531595569232 -->\n<g id=\"node3\" class=\"node\">\n<title>140531595569232</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"316,-140 227,-140 227,-121 316,-121 316,-140\"/>\n<text text-anchor=\"middle\" x=\"271.5\" y=\"-128\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n</g>\n<!-- 140531595569232&#45;&gt;140531595567168 -->\n<g id=\"edge1\" class=\"edge\">\n<title>140531595569232&#45;&gt;140531595567168</title>\n<path fill=\"none\" stroke=\"black\" d=\"M271.5,-120.75C271.5,-113.8 271.5,-103.85 271.5,-95.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"275,-95.09 271.5,-85.09 268,-95.09 275,-95.09\"/>\n</g>\n<!-- 140531595568800 -->\n<g id=\"node4\" class=\"node\">\n<title>140531595568800</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"281,-195 180,-195 180,-176 281,-176 281,-195\"/>\n<text text-anchor=\"middle\" x=\"230.5\" y=\"-183\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140531595568800&#45;&gt;140531595569232 -->\n<g id=\"edge2\" class=\"edge\">\n<title>140531595568800&#45;&gt;140531595569232</title>\n<path fill=\"none\" stroke=\"black\" d=\"M237.27,-175.75C243.06,-168.26 251.56,-157.28 258.64,-148.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"261.51,-150.14 264.86,-140.09 255.97,-145.86 261.51,-150.14\"/>\n</g>\n<!-- 140534660247264 -->\n<g id=\"node5\" class=\"node\">\n<title>140534660247264</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"135,-459 40,-459 40,-418 135,-418 135,-459\"/>\n<text text-anchor=\"middle\" x=\"87.5\" y=\"-447\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n<text text-anchor=\"middle\" x=\"87.5\" y=\"-436\" font-family=\"monospace\" font-size=\"10.00\">step1.fc.bias</text>\n<text text-anchor=\"middle\" x=\"87.5\" y=\"-425\" font-family=\"monospace\" font-size=\"10.00\">(1)</text>\n</g>\n<!-- 140534660247264&#45;&gt;140531595568800 -->\n<g id=\"edge3\" class=\"edge\">\n<title>140534660247264&#45;&gt;140531595568800</title>\n<path fill=\"none\" stroke=\"black\" d=\"M93.61,-417.94C106.18,-379.61 137.57,-292.99 182.5,-231 190.44,-220.04 201.39,-209.65 210.82,-201.66\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"213.06,-204.35 218.6,-195.32 208.64,-198.92 213.06,-204.35\"/>\n</g>\n<!-- 140534553595376 -->\n<g id=\"node6\" class=\"node\">\n<title>140534553595376</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"101,-849.5 0,-849.5 0,-830.5 101,-830.5 101,-849.5\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-837.5\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140534553595376&#45;&gt;140534660247264 -->\n<g id=\"edge4\" class=\"edge\">\n<title>140534553595376&#45;&gt;140534660247264</title>\n<path fill=\"none\" stroke=\"black\" d=\"M45.24,-830.23C34.59,-811.63 11.5,-766.55 11.5,-725.5 11.5,-725.5 11.5,-725.5 11.5,-558.5 11.5,-528.63 17.82,-520.42 33.5,-495 39.83,-484.73 48.36,-474.85 56.76,-466.33\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"59.33,-468.71 64.06,-459.23 54.45,-463.7 59.33,-468.71\"/>\n</g>\n<!-- 140534553592832 -->\n<g id=\"node13\" class=\"node\">\n<title>140534553592832</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"216,-789 115,-789 115,-770 216,-770 216,-789\"/>\n<text text-anchor=\"middle\" x=\"165.5\" y=\"-777\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140534553595376&#45;&gt;140534553592832 -->\n<g id=\"edge12\" class=\"edge\">\n<title>140534553595376&#45;&gt;140534553592832</title>\n<path fill=\"none\" stroke=\"black\" d=\"M67.47,-830.37C86.41,-820.73 117.29,-805.03 139.33,-793.81\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"141.14,-796.82 148.46,-789.17 137.96,-790.58 141.14,-796.82\"/>\n</g>\n<!-- 140534064448352 -->\n<g id=\"node7\" class=\"node\">\n<title>140534064448352</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"98,-921 3,-921 3,-891 98,-891 98,-921\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-909\" font-family=\"monospace\" font-size=\"10.00\">step0.fc.bias</text>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-898\" font-family=\"monospace\" font-size=\"10.00\">(1)</text>\n</g>\n<!-- 140534064448352&#45;&gt;140534553595376 -->\n<g id=\"edge5\" class=\"edge\">\n<title>140534064448352&#45;&gt;140534553595376</title>\n<path fill=\"none\" stroke=\"black\" d=\"M50.5,-890.8C50.5,-881.7 50.5,-869.79 50.5,-859.9\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54,-859.84 50.5,-849.84 47,-859.84 54,-859.84\"/>\n</g>\n<!-- 140534553595616 -->\n<g id=\"node8\" class=\"node\">\n<title>140534553595616</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"132,-514 43,-514 43,-495 132,-495 132,-514\"/>\n<text text-anchor=\"middle\" x=\"87.5\" y=\"-502\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140534553595616&#45;&gt;140534660247264 -->\n<g id=\"edge6\" class=\"edge\">\n<title>140534553595616&#45;&gt;140534660247264</title>\n<path fill=\"none\" stroke=\"black\" d=\"M87.5,-494.87C87.5,-488.22 87.5,-478.63 87.5,-469.28\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"91,-469.01 87.5,-459.01 84,-469.01 91,-469.01\"/>\n</g>\n<!-- 140534553594848 -->\n<g id=\"node9\" class=\"node\">\n<title>140534553594848</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"135,-569 40,-569 40,-550 135,-550 135,-569\"/>\n<text text-anchor=\"middle\" x=\"87.5\" y=\"-557\" font-family=\"monospace\" font-size=\"10.00\">ViewBackward0</text>\n</g>\n<!-- 140534553594848&#45;&gt;140534553595616 -->\n<g id=\"edge7\" class=\"edge\">\n<title>140534553594848&#45;&gt;140534553595616</title>\n<path fill=\"none\" stroke=\"black\" d=\"M87.5,-549.75C87.5,-542.8 87.5,-532.85 87.5,-524.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"91,-524.09 87.5,-514.09 84,-524.09 91,-524.09\"/>\n</g>\n<!-- 140534553594992 -->\n<g id=\"node10\" class=\"node\">\n<title>140534553594992</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"133,-624 44,-624 44,-605 133,-605 133,-624\"/>\n<text text-anchor=\"middle\" x=\"88.5\" y=\"-612\" font-family=\"monospace\" font-size=\"10.00\">SumBackward1</text>\n</g>\n<!-- 140534553594992&#45;&gt;140534553594848 -->\n<g id=\"edge8\" class=\"edge\">\n<title>140534553594992&#45;&gt;140534553594848</title>\n<path fill=\"none\" stroke=\"black\" d=\"M88.33,-604.75C88.2,-597.8 88.02,-587.85 87.85,-579.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"91.35,-579.02 87.66,-569.09 84.35,-579.15 91.35,-579.02\"/>\n</g>\n<!-- 140534553594800 -->\n<g id=\"node11\" class=\"node\">\n<title>140534553594800</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"236,-679 75,-679 75,-660 236,-660 236,-679\"/>\n<text text-anchor=\"middle\" x=\"155.5\" y=\"-667\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackwardBackward0</text>\n</g>\n<!-- 140534553594800&#45;&gt;140534553594992 -->\n<g id=\"edge9\" class=\"edge\">\n<title>140534553594800&#45;&gt;140534553594992</title>\n<path fill=\"none\" stroke=\"black\" d=\"M144.74,-659.98C134.73,-652.07 119.61,-640.11 107.57,-630.58\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"109.47,-627.63 99.46,-624.17 105.13,-633.12 109.47,-627.63\"/>\n</g>\n<!-- 140531595617904 -->\n<g id=\"node25\" class=\"node\">\n<title>140531595617904</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"230,-624 153,-624 153,-605 230,-605 230,-624\"/>\n<text text-anchor=\"middle\" x=\"191.5\" y=\"-612\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140534553594800&#45;&gt;140531595617904 -->\n<g id=\"edge26\" class=\"edge\">\n<title>140534553594800&#45;&gt;140531595617904</title>\n<path fill=\"none\" stroke=\"black\" d=\"M161.44,-659.75C166.48,-652.34 173.84,-641.5 180.01,-632.41\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"182.94,-634.33 185.67,-624.09 177.15,-630.39 182.94,-634.33\"/>\n</g>\n<!-- 140534553593072 -->\n<g id=\"node12\" class=\"node\">\n<title>140534553593072</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"210,-734 121,-734 121,-715 210,-715 210,-734\"/>\n<text text-anchor=\"middle\" x=\"165.5\" y=\"-722\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n</g>\n<!-- 140534553593072&#45;&gt;140534553594800 -->\n<g id=\"edge10\" class=\"edge\">\n<title>140534553593072&#45;&gt;140534553594800</title>\n<path fill=\"none\" stroke=\"black\" d=\"M163.85,-714.75C162.54,-707.8 160.66,-697.85 159.02,-689.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"162.41,-688.27 157.12,-679.09 155.54,-689.56 162.41,-688.27\"/>\n</g>\n<!-- 140534553592832&#45;&gt;140534553593072 -->\n<g id=\"edge11\" class=\"edge\">\n<title>140534553592832&#45;&gt;140534553593072</title>\n<path fill=\"none\" stroke=\"black\" d=\"M165.5,-769.75C165.5,-762.8 165.5,-752.85 165.5,-744.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"169,-744.09 165.5,-734.09 162,-744.09 169,-744.09\"/>\n</g>\n<!-- 140534553593456 -->\n<g id=\"node14\" class=\"node\">\n<title>140534553593456</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"215,-849.5 138,-849.5 138,-830.5 215,-830.5 215,-849.5\"/>\n<text text-anchor=\"middle\" x=\"176.5\" y=\"-837.5\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140534553593456&#45;&gt;140534553592832 -->\n<g id=\"edge13\" class=\"edge\">\n<title>140534553593456&#45;&gt;140534553592832</title>\n<path fill=\"none\" stroke=\"black\" d=\"M174.88,-830.37C173.33,-822.16 170.96,-809.54 168.99,-799.05\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"172.42,-798.35 167.13,-789.17 165.54,-799.64 172.42,-798.35\"/>\n</g>\n<!-- 140534553593888 -->\n<g id=\"node15\" class=\"node\">\n<title>140534553593888</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"246,-915.5 145,-915.5 145,-896.5 246,-896.5 246,-915.5\"/>\n<text text-anchor=\"middle\" x=\"195.5\" y=\"-903.5\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140534553593888&#45;&gt;140534553593456 -->\n<g id=\"edge14\" class=\"edge\">\n<title>140534553593888&#45;&gt;140534553593456</title>\n<path fill=\"none\" stroke=\"black\" d=\"M192.94,-896.37C190.15,-886.97 185.61,-871.67 182,-859.53\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"185.35,-858.5 179.14,-849.91 178.63,-860.49 185.35,-858.5\"/>\n</g>\n<!-- 140531595572368 -->\n<g id=\"node20\" class=\"node\">\n<title>140531595572368</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"284,-327 177,-327 177,-286 284,-286 284,-327\"/>\n<text text-anchor=\"middle\" x=\"230.5\" y=\"-315\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n<text text-anchor=\"middle\" x=\"230.5\" y=\"-304\" font-family=\"monospace\" font-size=\"10.00\">step1.fc.weight</text>\n<text text-anchor=\"middle\" x=\"230.5\" y=\"-293\" font-family=\"monospace\" font-size=\"10.00\">(1, 5)</text>\n</g>\n<!-- 140534553593888&#45;&gt;140531595572368 -->\n<g id=\"edge20\" class=\"edge\">\n<title>140534553593888&#45;&gt;140531595572368</title>\n<path fill=\"none\" stroke=\"black\" d=\"M204.52,-896.26C223.22,-877.15 264.5,-829.43 264.5,-780.5 264.5,-780.5 264.5,-780.5 264.5,-437.5 264.5,-402.08 252.54,-362.94 242.8,-336.96\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"245.96,-335.4 239.08,-327.34 239.43,-337.93 245.96,-335.4\"/>\n</g>\n<!-- 140531595612944 -->\n<g id=\"node16\" class=\"node\">\n<title>140531595612944</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"249,-987 142,-987 142,-957 249,-957 249,-987\"/>\n<text text-anchor=\"middle\" x=\"195.5\" y=\"-975\" font-family=\"monospace\" font-size=\"10.00\">step0.fc.weight</text>\n<text text-anchor=\"middle\" x=\"195.5\" y=\"-964\" font-family=\"monospace\" font-size=\"10.00\">(1, 5)</text>\n</g>\n<!-- 140531595612944&#45;&gt;140534553593888 -->\n<g id=\"edge15\" class=\"edge\">\n<title>140531595612944&#45;&gt;140534553593888</title>\n<path fill=\"none\" stroke=\"black\" d=\"M195.5,-956.8C195.5,-947.7 195.5,-935.79 195.5,-925.9\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"199,-925.84 195.5,-915.84 192,-925.84 199,-925.84\"/>\n</g>\n<!-- 140531595567888 -->\n<g id=\"node17\" class=\"node\">\n<title>140531595567888</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"394,-789 293,-789 293,-770 394,-770 394,-789\"/>\n<text text-anchor=\"middle\" x=\"343.5\" y=\"-777\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140531595567888&#45;&gt;140531595569232 -->\n<g id=\"edge27\" class=\"edge\">\n<title>140531595567888&#45;&gt;140531595569232</title>\n<path fill=\"none\" stroke=\"black\" d=\"M340.67,-769.93C334.94,-751.8 322.5,-708.13 322.5,-670.5 322.5,-670.5 322.5,-670.5 322.5,-239.5 322.5,-204.28 300.05,-168.25 284.94,-147.99\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"287.7,-145.84 278.81,-140.09 282.17,-150.13 287.7,-145.84\"/>\n</g>\n<!-- 140531595567888&#45;&gt;140534553593072 -->\n<g id=\"edge16\" class=\"edge\">\n<title>140531595567888&#45;&gt;140534553593072</title>\n<path fill=\"none\" stroke=\"black\" d=\"M314.9,-769.98C284.89,-761.05 237.52,-746.94 204.18,-737.02\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"204.91,-733.58 194.32,-734.08 202.91,-740.29 204.91,-733.58\"/>\n</g>\n<!-- 140531595613184 -->\n<g id=\"node18\" class=\"node\">\n<title>140531595613184</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"382,-855 305,-855 305,-825 382,-825 382,-855\"/>\n<text text-anchor=\"middle\" x=\"343.5\" y=\"-843\" font-family=\"monospace\" font-size=\"10.00\">meta_param</text>\n<text text-anchor=\"middle\" x=\"343.5\" y=\"-832\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140531595613184&#45;&gt;140531595567888 -->\n<g id=\"edge17\" class=\"edge\">\n<title>140531595613184&#45;&gt;140531595567888</title>\n<path fill=\"none\" stroke=\"black\" d=\"M343.5,-824.84C343.5,-817.21 343.5,-807.7 343.5,-799.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"347,-799.27 343.5,-789.27 340,-799.27 347,-799.27\"/>\n</g>\n<!-- 140534553594272 -->\n<g id=\"node19\" class=\"node\">\n<title>140534553594272</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"269,-250 192,-250 192,-231 269,-231 269,-250\"/>\n<text text-anchor=\"middle\" x=\"230.5\" y=\"-238\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140534553594272&#45;&gt;140531595568800 -->\n<g id=\"edge18\" class=\"edge\">\n<title>140534553594272&#45;&gt;140531595568800</title>\n<path fill=\"none\" stroke=\"black\" d=\"M230.5,-230.75C230.5,-223.8 230.5,-213.85 230.5,-205.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"234,-205.09 230.5,-195.09 227,-205.09 234,-205.09\"/>\n</g>\n<!-- 140531595572368&#45;&gt;140534553594272 -->\n<g id=\"edge19\" class=\"edge\">\n<title>140531595572368&#45;&gt;140534553594272</title>\n<path fill=\"none\" stroke=\"black\" d=\"M230.5,-285.95C230.5,-277.85 230.5,-268.5 230.5,-260.47\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"234,-260.26 230.5,-250.26 227,-260.26 234,-260.26\"/>\n</g>\n<!-- 140534553593504 -->\n<g id=\"node21\" class=\"node\">\n<title>140534553593504</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"236,-382 147,-382 147,-363 236,-363 236,-382\"/>\n<text text-anchor=\"middle\" x=\"191.5\" y=\"-370\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140534553593504&#45;&gt;140531595572368 -->\n<g id=\"edge21\" class=\"edge\">\n<title>140534553593504&#45;&gt;140531595572368</title>\n<path fill=\"none\" stroke=\"black\" d=\"M196.76,-362.87C201.03,-355.87 207.29,-345.59 213.27,-335.78\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"216.4,-337.37 218.61,-327.01 210.42,-333.73 216.4,-337.37\"/>\n</g>\n<!-- 140534553592976 -->\n<g id=\"node22\" class=\"node\">\n<title>140534553592976</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"231,-448 154,-448 154,-429 231,-429 231,-448\"/>\n<text text-anchor=\"middle\" x=\"192.5\" y=\"-436\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140534553592976&#45;&gt;140534553593504 -->\n<g id=\"edge22\" class=\"edge\">\n<title>140534553592976&#45;&gt;140534553593504</title>\n<path fill=\"none\" stroke=\"black\" d=\"M192.37,-428.87C192.22,-419.66 191.99,-404.79 191.8,-392.77\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"195.3,-392.35 191.64,-382.41 188.3,-392.46 195.3,-392.35\"/>\n</g>\n<!-- 140534553593216 -->\n<g id=\"node23\" class=\"node\">\n<title>140534553593216</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"232,-514 155,-514 155,-495 232,-495 232,-514\"/>\n<text text-anchor=\"middle\" x=\"193.5\" y=\"-502\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140534553593216&#45;&gt;140534553592976 -->\n<g id=\"edge23\" class=\"edge\">\n<title>140534553593216&#45;&gt;140534553592976</title>\n<path fill=\"none\" stroke=\"black\" d=\"M193.37,-494.87C193.22,-485.66 192.99,-470.79 192.8,-458.77\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"196.3,-458.35 192.64,-448.41 189.3,-458.46 196.3,-458.35\"/>\n</g>\n<!-- 140534553593552 -->\n<g id=\"node24\" class=\"node\">\n<title>140534553593552</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"236,-569 153,-569 153,-550 236,-550 236,-569\"/>\n<text text-anchor=\"middle\" x=\"194.5\" y=\"-557\" font-family=\"monospace\" font-size=\"10.00\">MmBackward0</text>\n</g>\n<!-- 140534553593552&#45;&gt;140534553593216 -->\n<g id=\"edge24\" class=\"edge\">\n<title>140534553593552&#45;&gt;140534553593216</title>\n<path fill=\"none\" stroke=\"black\" d=\"M194.33,-549.75C194.2,-542.8 194.02,-532.85 193.85,-524.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"197.35,-524.02 193.66,-514.09 190.35,-524.15 197.35,-524.02\"/>\n</g>\n<!-- 140531595617904&#45;&gt;140534553593552 -->\n<g id=\"edge25\" class=\"edge\">\n<title>140531595617904&#45;&gt;140534553593552</title>\n<path fill=\"none\" stroke=\"black\" d=\"M192,-604.75C192.39,-597.8 192.95,-587.85 193.45,-579.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"196.94,-579.27 194.01,-569.09 189.95,-578.88 196.94,-579.27\"/>\n</g>\n</g>\n</svg>\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "class MetaNet(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Linear(dim, 1, bias=True)\n",
    "\n",
    "    def forward(self, x, meta_param):\n",
    "        return self.fc(x) + meta_param\n",
    "\n",
    "\n",
    "dim = 5\n",
    "batch_size = 2\n",
    "net = MetaNet(dim)\n",
    "\n",
    "xs = torch.ones((batch_size, dim))\n",
    "ys = torch.ones((batch_size, 1))\n",
    "\n",
    "optimizer = torchopt.MetaSGD(net, lr=1e-3)\n",
    "meta_param = torch.tensor(1.0, requires_grad=True)\n",
    "\n",
    "# Set enable_visual\n",
    "net_state_0 = torchopt.extract_state_dict(net, enable_visual=True, visual_prefix='step0.')\n",
    "\n",
    "pred = net(xs, meta_param)\n",
    "loss = F.mse_loss(pred, ys)\n",
    "optimizer.step(loss)\n",
    "\n",
    "# Set enable_visual\n",
    "net_state_1 = torchopt.extract_state_dict(net, enable_visual=True, visual_prefix='step1.')\n",
    "\n",
    "pred = net(xs, meta_param)\n",
    "loss = F.mse_loss(pred, torch.ones_like(pred))\n",
    "\n",
    "# Draw computation graph\n",
    "display(\n",
    "    torchopt.visual.make_dot(\n",
    "        loss, [net_state_0, net_state_1, {'meta_param': meta_param, 'loss': loss}]\n",
    "    )\n",
    ")"
   ]
  }
 ],
 "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.9.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "2a8cc1ff2cbc47027bf9993941710d9ab9175f14080903d9c7c432ee63d681da"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}