解决 HiddenLayer 可视化不显示神经网络数据流维度大小

解决 HiddenLayer 可视化不显示网络数据流维度大小

现象：

使用 HiddenLayer 可视化神经网络结构的时候，
不显示中间数据的大小

解决方法：修改底层源码

找到 pytorch_builder.py 文件，修改 get_shape(torch_node) 函数如下：

def get_shape(torch_node):
    """Return the output shape of the given Pytorch node."""
    # Extract node output shape from the node string representation
    # This is a hack because there doesn't seem to be an official way to do it.
    # See my quesiton in the PyTorch forum:
    # https://discuss.pytorch.org/t/node-output-shape-from-trace-graph/24351/2
    # TODO: find a better way to extract output shape
    # TODO: Assuming the node has one output. Update if we encounter a multi-output node.
    shape = torch_node.output().type().sizes()
    return shape

解决后的效果：

能显示中间过程的数据大小：

参考: 参考链接