Visualization Tools that can show proper branch and merge in Graphs

Question

I want help finding the visualization tool can draw similar architecture as given in the image below

Keras visualization produces something similar to this graph. But I'm working in Pytorch. I tried using Netron and Graphviz, both produce graphs that do not show the branching and merging properly.

This is the plot I rendered in Graphviz for another branching and merging architecture.

Here's the python code I used to render the plot

graph {
    graph [rankdir=RL]
    input [label=Input]
    conv1 [label="Conv 1"]
    conv2 [label="Conv 2"]
    conv3 [label="Conv 3"]
    conv4 [label="Conv 4"]
    conv5 [label="Conv 5"]
    conv6 [label="Conv 6"]
    conv7 [label="Conv 7"]
    conv8 [label="Conv 8"]
    conv9 [label="Conv 9"]
    batch1 [label="Batch Norm1"]
    batch2 [label="Batch Norm2"]
    batch3 [label="Batch Norm3"]
    caps1_a [label="Caps 1(a)"]
    caps1_b [label="Caps 1(b)"]
    caps2_a [label="Caps 2(a)"]
    caps2_b [label="Caps 2(b)"]
    caps3_a [label="Caps 3(a)"]
    caps3_b [label="Caps 3(b)"]
    sum1 [label=Sum]
    sum2 [label=Sum]
    sum3 [label=Sum]
    stack [label=Stack]
    sum4 [label=Sum]
    softmax [label=Softmax]
    input -- conv1 [constraint=true]
    conv1 -- conv2 [constraint=false]
    conv2 -- conv3 [constraint=false]
    conv3 -- conv4 [constraint=false]
    conv4 -- conv5 [constraint=false]
    conv5 -- conv6 [constraint=false]
    conv6 -- conv7 [constraint=false]
    conv7 -- conv8 [constraint=false]
    conv8 -- conv9 [constraint=false]
    conv3 -- batch1 [constraint=false]
    conv6 -- batch2 [constraint=false]
    conv9 -- batch3 [constraint=false]
    batch1 -- caps1_a [constraint=false]
    batch2 -- caps2_a [constraint=false]
    batch3 -- caps3_a [constraint=false]
    caps1_a -- caps1_b [constraint=false]
    caps2_a -- caps2_b [constraint=false]
    caps3_a -- caps3_b [constraint=false]
    caps1_b -- sum1 [constraint=false]
    caps2_b -- sum2 [constraint=false]
    caps3_b -- sum3 [constraint=false]
    sum1 -- stack [label=w1 constraint=false]
    sum2 -- stack [label=w2 constraint=false]
    sum3 -- stack [label=w3 constraint=false]
    stack -- sum4 [constraint=false]
    sum4 -- softmax [constraint=false]
    input -- conv2 [label=k1 constraint=false]
    conv2 -- conv4 [label=k2 constraint=false]
    conv4 -- conv6 [label=k3 constraint=false]
    conv6 -- conv8 [label=k4 constraint=false]
}

Answer 1

draw.io is a great visualization tool, it helped me draw diagrams as shown in the question.

Answer 2

You could use Tensorboard, via the tensorboardX interface.

That allows you to load models from PyTorch (and Chainer, MXNet etc.) into Tensorboard. This will then show you the full graph interactively.

From the homepage: