How to draw graph in tensorboardX with a two inputs model?

The inputs of model are like this

output = model(input1, input2)

How can I use function writer.add_graph(model, input) ?

Here is an example:

import torch
import torch.nn as nn
from tensorboardX import SummaryWriter

input1 = torch.ones(1, 2)
input2 = torch.randn(1, 2)
batch_input = torch.cat((input1, input2), 0) # Two input

class SimpleLinear(nn.Module):

def __init__(self):
    super(SimpleLinear, self).__init__()
    self.slin = nn.Linear(2, 4)

def forward(self, x):
    return self.slin(x)

with SummaryWriter(‘runs/graph2’, comment=‘SimpleLinear’) as w:

w.add_graph(SimpleLinear(), batch_input, True)

Hi,

you can do:

writer.add_graph(model, (input1, input2) )

Here is an example

2 Likes

hi I got the error: Type ‘Tuple[Tensor, Dict[str, List[int]]]’ cannot be traced. Only Tensors and (possibly nested) Lists, Dicts, and Tuples of Tensors can be traced
The code is below:
writer.add_graph(flow_net, (torch.DoubleTensor(dataset.datas_vali[:2,:]), idx))