I used “torch.fx.symbolic_trace” to trace the pretrained Alexnet, and then I got a graph of Alexnet. But when I print the node of graph, I only got “feature_0. feature_1…”. But actually I want “conv, relu, maxpooling…”
Here is my code:
traced = symbolic_trace(module)
graph = traced.graph
for node in graph.nodes:
print(node.name)
Ang my results:
x
features_0
features_1
features_2
features_3
features_4
features_5
features_6
features_7
features_8
features_9
features_10
features_11
features_12
avgpool
flatten_1
classifier_0
classifier_1
classifier_2
classifier_3
classifier_4
classifier_5
classifier_6
output
I want to print op name of each node like:
Conv
Relu
MaxPool
Conv
Relu
MaxPool
Conv
Relu
Conv
Relu
Conv
Relu
MaxPool
AveragePool
Flatten
Dropout
Linear
Relu
Dropout
Linear
Relu
Linear
I found the when the model based on nn.Sequential and it cannot trace the op name. Could anyone help me with this problem?