Using the following code — which comes from the PyTorch demo notebook — I get the output in the subsequent image. The Pytorch demo notebook suggests in its first example that tensor shapes are printed on the edges between nodes. However, I have found this not to be the case in my attempts. How can I get the tensor shapes printed along the edges? Is this an issue, or is there documentation I can refer to fix the output?
(By the way, I am running python3.8 with the latest version of Pytorch (1.6.0) and hiddenlayer).
import torchvision
import hiddenlayer as hl
model = torchvision.models.resnet101()
# Rather than using the default transforms, build custom ones to group
# nodes of residual and bottleneck blocks.
transforms = [
# Fold Conv, BN, RELU layers into one
hl.transforms.Fold("Conv > BatchNorm > Relu", "ConvBnRelu"),
# Fold Conv, BN layers together
hl.transforms.Fold("Conv > BatchNorm", "ConvBn"),
# Fold bottleneck blocks
hl.transforms.Fold(
"""
((ConvBnRelu > ConvBnRelu > ConvBn) | ConvBn) > Add > Relu
""",
"BottleneckBlock",
"Bottleneck Block",
),
# Fold residual blocks
hl.transforms.Fold(
"""ConvBnRelu > ConvBnRelu > ConvBn > Add > Relu""",
"ResBlock",
"Residual Block",
),
# Fold repeated blocks
hl.transforms.FoldDuplicates(),
]
# Display graph using the transforms above
res_graph = hl.build_graph(
model, torch.zeros([1, 3, 224, 224]), transforms=transforms
)
res_graph.save(path="./log/res_graph.pdf")
Using the following code — which comes from the PyTorch demo notebook — I get the output in the subsequent image. The Pytorch demo notebook suggests in its first example that tensor shapes are printed on the edges between nodes. However, I have found this not to be the case in my attempts. How can I get the tensor shapes printed along the edges? Is this an issue, or is there documentation I can refer to fix the output?
(By the way, I am running python3.8 with the latest version of Pytorch (1.6.0) and hiddenlayer).