corrected tree height

Author: sethjuarez

graph.ipynb

@@ -1 +1 @@
-{"cells":[{"cell_type":"code","execution_count":1,"metadata":{},"outputs":[],"source":"import re\nimport json\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport utils.viz as torchviz\ntorchviz.init()"},{"cell_type":"code","execution_count":2,"metadata":{},"outputs":[],"source":"x = torch.tensor([[255., 232, 132, 12, 343, 21, 0, 32, 12]]) / 255\ny = torch.tensor([1., 0, 0])"},{"cell_type":"code","execution_count":3,"metadata":{},"outputs":[],"source":"w = torch.rand(9, 3, requires_grad=True)\nb = torch.rand(1, 3, requires_grad=True)"},{"cell_type":"code","execution_count":4,"metadata":{},"outputs":[],"source":"h = x.mm(w) + b\nloss = ((h - y)**2).mean()\n\n# here\n"},{"cell_type":"code","execution_count":5,"metadata":{},"outputs":[],"source":"c = torchviz.draw(loss)\nc"},{"cell_type":"code","execution_count":9,"metadata":{},"outputs":[],"source":"loss.backward()"},{"cell_type":"code","execution_count":10,"metadata":{},"outputs":[],"source":"print(w.grad, '\\n', b.grad)"},{"cell_type":"markdown","metadata":{},"outputs":[],"source":["# Linear"]},{"cell_type":"code","execution_count":0,"metadata":{},"outputs":[],"source":"class LinearModel(nn.Module):\n    def __init__(self, x, y):\n        super(LinearModel, self).__init__()\n        self.layer1 = nn.Linear(x, y)\n\n    def forward(self, x):\n        x = self.layer1(x)\n        return F.softmax(x, dim=1)"},{"cell_type":"code","execution_count":0,"metadata":{},"outputs":[],"source":"model = LinearModel(9, 3)\ncost = nn.MSELoss()\npred = model(x)\nloss = cost(pred, y)\ntorchviz.draw(loss)"},{"cell_type":"markdown","metadata":{},"outputs":[],"source":["# Neural Networks"]},{"cell_type":"code","execution_count":0,"metadata":{},"outputs":[],"source":"class NeuralNework(nn.Module):\n    def __init__(self, x, y):\n        super(NeuralNework, self).__init__()\n        hidden = int(x/2)\n        self.layer1 = nn.Linear(x, hidden)\n        self.layer2 = nn.Linear(hidden, hidden)\n        self.output = nn.Linear(hidden, y)\n\n    def forward(self, x):\n        x = F.relu(self.layer1(x))\n        x = F.relu(self.layer2(x))\n        x = self.output(x)\n        return F.softmax(x, dim=1)"},{"cell_type":"code","execution_count":0,"metadata":{},"outputs":[],"source":"model = NeuralNework(9, 3)\ncost = nn.MSELoss()\npred = model(x)\nloss = cost(pred, y)\ntorchviz.draw(loss)"},{"cell_type":"markdown","metadata":{},"outputs":[],"source":["# Convolutional Neural Network"]},{"cell_type":"code","execution_count":0,"metadata":{},"outputs":[],"source":"from torch.utils.data import DataLoader\nfrom torch.utils.data.dataset import Dataset\nfrom torchvision import datasets, transforms\n\ndigits = datasets.MNIST('data', train=True, download=True,\n                    transform=transforms.Compose([\n                        transforms.ToTensor(),\n                        transforms.Lambda(lambda x: x.view(28*28))\n                    ]),\n                    target_transform=transforms.Compose([\n                        transforms.Lambda(lambda y: \n                                          torch.zeros(10, dtype=torch.float).scatter_(0, torch.tensor(y), value=1))\n                    ])\n                 )"},{"cell_type":"code","execution_count":0,"metadata":{},"outputs":[],"source":"class CNN(nn.Module):\n    def __init__(self):\n        super(CNN, self).__init__()\n        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n        self.conv2_drop = nn.Dropout2d()\n        self.fc1 = nn.Linear(320, 50)\n        self.fc2 = nn.Linear(50, 10)\n\n    def forward(self, x):\n        x = x.view(-1, 1, 28, 28)\n        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n        x = x.view(-1, 320)\n        x = F.relu(self.fc1(x))\n        x = F.dropout(x, training=self.training)\n        x = self.fc2(x)\n        return F.softmax(x, dim=1)"},{"cell_type":"code","execution_count":0,"metadata":{},"outputs":[],"source":"model = CNN()\ncost = torch.nn.BCELoss()\npred = model(digits[0][0])\nloss = cost(pred, digits[0][1].view(1, -1))\ntorchviz.draw(loss)"}],"nbformat":4,"nbformat_minor":2,"metadata":{"language_info":{"name":"python","codemirror_mode":{"name":"ipython","version":3}},"orig_nbformat":2,"file_extension":".py","mimetype":"text/x-python","name":"python","npconvert_exporter":"python","pygments_lexer":"ipython3","version":3}}
+{"cells":[{"cell_type":"code","execution_count":1,"metadata":{},"outputs":[],"source":"import re\nimport json\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport utils.viz as torchviz\ntorchviz.init()"},{"cell_type":"code","execution_count":2,"metadata":{},"outputs":[],"source":"x = torch.tensor([[255., 232, 132, 12, 343, 21, 0, 32, 12]]) / 255\ny = torch.tensor([1., 0, 0])"},{"cell_type":"code","execution_count":3,"metadata":{},"outputs":[],"source":"w = torch.rand(9, 3, requires_grad=True)\nb = torch.rand(1, 3, requires_grad=True)"},{"cell_type":"code","execution_count":4,"metadata":{},"outputs":[],"source":"h = x.mm(w) + b\nloss = ((h - y)**2).mean()\n\n# here\n"},{"cell_type":"code","execution_count":5,"metadata":{},"outputs":[],"source":"c = torchviz.draw(loss)\nc"},{"cell_type":"code","execution_count":6,"metadata":{},"outputs":[],"source":"loss.backward()"},{"cell_type":"code","execution_count":7,"metadata":{},"outputs":[],"source":"print(w.grad, '\\n', b.grad)"},{"cell_type":"markdown","metadata":{},"outputs":[],"source":"# Linear"},{"cell_type":"code","execution_count":8,"metadata":{},"outputs":[],"source":"class LinearModel(nn.Module):\n    def __init__(self, x, y):\n        super(LinearModel, self).__init__()\n        self.layer1 = nn.Linear(x, y)\n\n    def forward(self, x):\n        x = self.layer1(x)\n        return F.softmax(x, dim=1)"},{"cell_type":"code","execution_count":9,"metadata":{},"outputs":[],"source":"model = LinearModel(9, 3)\ncost = nn.MSELoss()\npred = model(x)\nloss = cost(pred, y.view(1, -1))\ntorchviz.draw(loss)"},{"cell_type":"markdown","metadata":{},"outputs":[],"source":"# Neural Networks"},{"cell_type":"code","execution_count":10,"metadata":{},"outputs":[],"source":"class NeuralNework(nn.Module):\n    def __init__(self, x, y):\n        super(NeuralNework, self).__init__()\n        hidden = int(x/2)\n        self.layer1 = nn.Linear(x, hidden)\n        self.layer2 = nn.Linear(hidden, hidden)\n        self.output = nn.Linear(hidden, y)\n\n    def forward(self, x):\n        x = F.relu(self.layer1(x))\n        x = F.relu(self.layer2(x))\n        x = self.output(x)\n        return F.softmax(x, dim=1)"},{"cell_type":"code","execution_count":11,"metadata":{},"outputs":[],"source":"model = NeuralNework(9, 3)\ncost = nn.MSELoss()\npred = model(x)\nloss = cost(pred, y.view(1, -1))\ntorchviz.draw(loss)"},{"cell_type":"markdown","metadata":{},"outputs":[],"source":"# Convolutional Neural Network"},{"cell_type":"code","execution_count":12,"metadata":{},"outputs":[],"source":"from torch.utils.data import DataLoader\nfrom torch.utils.data.dataset import Dataset\nfrom torchvision import datasets, transforms\n\ndigits = datasets.MNIST('data', train=True, download=True,\n                    transform=transforms.Compose([\n                        transforms.ToTensor(),\n                        transforms.Lambda(lambda x: x.view(28*28))\n                    ]),\n                    target_transform=transforms.Compose([\n                        transforms.Lambda(lambda y: \n                                          torch.zeros(10, dtype=torch.float).scatter_(0, torch.tensor(y), value=1))\n                    ])\n                 )"},{"cell_type":"code","execution_count":13,"metadata":{},"outputs":[],"source":"class CNN(nn.Module):\n    def __init__(self):\n        super(CNN, self).__init__()\n        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n        self.conv2_drop = nn.Dropout2d()\n        self.fc1 = nn.Linear(320, 50)\n        self.fc2 = nn.Linear(50, 10)\n\n    def forward(self, x):\n        x = x.view(-1, 1, 28, 28)\n        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n        x = x.view(-1, 320)\n        x = F.relu(self.fc1(x))\n        x = F.dropout(x, training=self.training)\n        x = self.fc2(x)\n        return F.softmax(x, dim=1)"},{"cell_type":"code","execution_count":14,"metadata":{},"outputs":[],"source":"model = CNN()\ncost = torch.nn.BCELoss()\npred = model(digits[0][0])\nloss = cost(pred, digits[0][1].view(1, -1))\ntorchviz.draw(loss)"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":""}],"nbformat":4,"nbformat_minor":2,"metadata":{"language_info":{"name":"python","codemirror_mode":{"name":"ipython","version":3}},"orig_nbformat":2,"file_extension":".py","mimetype":"text/x-python","name":"python","npconvert_exporter":"python","pygments_lexer":"ipython3","version":3}}

utils/viz.py

@@ -8,19 +8,19 @@
 def init():
     return HTML('<script src="https://cdn.jsdelivr.net/npm/vega@5"></script><span>torchviz initialized!</span>')
 
-def get_dict(g, width=650, height=300):
+def get_dict(g, depth):
     d = {
         "$schema": "https://vega.github.io/schema/vega/v5.json",
-        "width": width,
-        "height": height,
+        "width": 650,
+        "height": depth*30,
         "padding": 5,    
         "signals": [
             {
                 "name": "labels", "value": True,
                 "bind": {"input": "checkbox"}
             },
             { 
-                "name": "method", "value": "cluster",
+                "name": "method", "value": "tidy",
                 "bind": {"input": "select", "options": ["tidy", "cluster"]} 
             },
             { 
@@ -110,7 +110,7 @@ def get_dict(g, width=650, height=300):
     d['data'][0]['values'] = g
     return d
 
-def build_graph(g, elements=[], parentId=-1):
+def build_graph(g, elements=[], parentId=-1, depth=0):
     elm = { 'id': len(elements), 'parentId': None if parentId==-1 else parentId}
     if g == None:
         elm['name'] = 'Const'
@@ -127,10 +127,11 @@ def build_graph(g, elements=[], parentId=-1):
     elements.append(elm)
 
     if g != None and g.next_functions != None:
+        depth = depth+1
         for subg in g.next_functions:
-            build_graph(subg[0], elements, elm['id'])
+            elements, depth = build_graph(subg[0], elements, elm['id'], depth)
 
-    return elements
+    return elements, depth
 
 def draw_graph(graph):
     return Javascript("""
@@ -145,9 +146,6 @@ def draw_graph(graph):
 """ % json.dumps(graph))
 
 def draw(g):
-    graph = build_graph(g.grad_fn, elements=[])
-    j = get_dict(graph)
-    #print(json.dumps(j, indent=2))
-    #return alt.Chart.from_dict(j)
-    #return json.dumps(j)
+    graph, depth = build_graph(g.grad_fn, elements=[])
+    j = get_dict(graph, depth)
     return draw_graph(j)