[Work in progress] visualization of autograd graph

lanpa · July 18, 2017, 12:40pm

http://35.197.26.245:6006/#graphs

I am wondering whether the route created wrt to the following code is correct, especially the two ConvNd nodes below the MaxPool layer. Any ideas?

ps. the graph is traced according to the Varible.grad_fn attribute recursively.

github.com

lanpa/tensorboard-pytorch/blob/master/demo_graph.py

import torch
import torch.nn as nn
import torchvision.utils as vutils
import numpy as np
import torch.nn.functional as F
import torchvision.models as models
from tensorboardX import SummaryWriter

class Mnist(nn.Module):
    def __init__(self):
        super(Mnist, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)
        self.bn = nn.BatchNorm2d(20)
    def forward(self, x):
        x = F.max_pool2d(self.conv1(x), 2)
        x = F.relu(x)+F.relu(-x)

This file has been truncated. show original