Visualise Network Architecture

1

I am trying to visualise my NN architecture. I use this code : PINNs/Burgers Inference (PyTorch).ipynb at master · jayroxis/PINNs (github.com) which is available on Github.

This is my NN:

class DNN(torch.nn.Module):

def __init__(self, layers):
    super(DNN, self).__init__()
    
    # parameters
    self.depth = len(layers) - 1
    
    # set up layer order dict
    self.activation = torch.nn.Tanh
    
    layer_list = list()
    for i in range(self.depth - 1): 
        layer_list.append(
            ('layer_%d' % i, torch.nn.Linear(layers[i], layers[i+1]))
        )
        layer_list.append(('activation_%d' % i, self.activation()))
        
    layer_list.append(
        ('layer_%d' % (self.depth - 1), torch.nn.Linear(layers[-2], layers[-1]))
    )
    layerDict = OrderedDict(layer_list)
    
    # deploy layers
    self.layers = torch.nn.Sequential(layerDict)
    
def forward(self, x):
    out = self.layers(x)
    return out

How can I visualise it, best as a PNG file? I would like to see the layers, neurons, weights biases etc.

Thanks in advance :slight_smile:

2

Refer this