Visualization of pytorch model based on netron

Visualization effect chart

example

from models.seg_hrnet import get_seg_model
import yaml
import torch
import os

if __name__ == '__main__':
    cfgPath = './experiments/cityscapes/seg_hrnet_w18_small_v1_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml'
    with open(cfgPath, 'r') as f:
        data = yaml.load(f)

    model = get_seg_model(data)
    img = torch.randn((2, 3, 512, 512))
    output = model(img)
    print(output.size())