Tensorboardx: A few questions about add_graph

[JiamingSuen]

Hi @lanpa , thanks for this amazing tool. I'm trying to use add_graph in my own project, where I met these questions:

1. How to use add_graph in the training iteration? Am I supposed to initialize the SummaryWriter outside the main training loop and call w.add_graph inside the loop? If so how do I get the model output? Or this feature is only possible for dummy input and being called only once? Why it's necessary to feed input to the model? It would be nice if you can provide an example with add_graph used in a complete project(instead of demo_graph.py).
2. Does it support custom CUDA layer? I tried with a custom correlation layer, and no matter what kind of data is fed into the model(random or image) the custom correlation failed for unknown reason.

Thanks for your time!

[lanpa]

You are welcome. I answer them one by one:

1. The overhead of graph visualization is huge. It's better to put outside your training loop.
2. The dummy input is used by PyTorch's JIT forwarding, whose by-product can be used to build the TB graph.
3. I use graph visualization to check the correctness of a model or teaching in the class. If you have idea on better example, I will try to implement.

1. Possibly yes (after some hard work). What are the error messages?

[JiamingSuen]

Thanks for your reply! Now I understand the purpose of this feature, which to me is not identical with the usage in TensorFlow.
For your question, the error message itself is not so helpful, it just says it cannot traceback into legacy code(which is the correlation CUDA kernel). The error only happens when I feed data with SummaryWriter so I believe it's related. I can reproduce it if needed but I don't think it will provide much information. Code defining this layer is as follows, the number of input/output channels is not clearly defined in this line of code. Is it possible for tensorboardX to show this layer on the graph?

self.corr = Correlation(pad_size=20, kernel_size=1, max_displacement=20, stride1=1, stride2=2, corr_multiply=1)

This correlation layer implementation is from here.

[PengLU1101]

Could you give an example to use add_graph_onnx()?

[lanpa]

@JiamingSuen There is a hacky solution. Since you only need to draw once, and result is don't care. You can replace the GPU module with the following.

class Correlation(nn.Module):
    def __init__(self):
        super(Correlation, self).__init__()

    def forward(self, x):
        # implement a fake correlation 
        #(as long as the output size is the same as real correlation) 
        return 

I will give a shot on the real GPU code in the weekend.

[lanpa]

@PengLU1101 The onnx code is here:
https://github.com/lanpa/tensorboard-pytorch/blob/9a3f9cca7057af515df64fdbfa523dbf16c1211d/onnx_graph.py
I hided it since onnx lacks scope information now, the resulting graph is almost useless.

[kunguang]

如果我的graph里有embedding层，是不是graph就读取不了

[lanpa]

"Will it work if there is an embedding layer in the graph?"

Haven't test that. Do you have any error message?

[kunguang]

def parse(graph):
scope = {}
for n in graph.nodes():
inputs = [i.uniqueName() for i in n.inputs()]
for i in range(1, len(inputs)):
scope[inputs[i]] = n.scopeName()
for i in range(1, len(inputs)): 改成 for i in range(0, len(inputs))，就可以了~