Python中TorchScript模型无法获取add等运算层的问题

为什么named_modules()看不到add层？ #

model.named_modules()只会遍历模型中的子模块实例（比如Conv2d、Bottleneck这类继承自torch.nn.Module的类），而add这类张量运算属于算子操作，是直接写在模块forward方法里的代码逻辑，并非独立的Module子类实例，所以不会被named_modules()捕获到。

你在Netron里看到的add层，对应TorchScript模型forward方法里的张量相加逻辑，属于计算图中的运算节点，而非子模块。

如何查看这些add层的信息？ #

1. 打印模型的forward脚本 #

TorchScript模型可通过model.code属性查看编译后的forward方法代码，里面会包含所有运算逻辑（包括add操作）：

model = torch.jit.load('best5.torchscript')
print(model.code)

以你示例中的C3模块为例，其forward代码里会存在类似x = x + y或torch.add(x, y)的逻辑，这就是Netron中add层对应的代码实现。

2. 遍历TorchScript计算图节点 #

通过获取模型的计算图，可以直接查看所有运算节点，包括add这类操作：

model = torch.jit.load('best5.torchscript')
graph = model.forward.graph
for node in graph.nodes():
    print(node.kind(), node.outputs())

其中node.kind()会显示节点类型，比如aten::add就是对应add运算的节点。

3. 导出为ONNX后解析 #

将TorchScript模型导出为ONNX格式，ONNX结构会包含所有运算节点，可通过onnx库解析查看：

import torch
import onnx

model = torch.jit.load('best5.torchscript')
dummy_input = torch.randn(1, 3, 640, 640)  # 需匹配模型实际输入尺寸
torch.onnx.export(model, dummy_input, 'model.onnx')

# 解析ONNX模型中的Add节点
onnx_model = onnx.load('model.onnx')
for node in onnx_model.graph.node:
    if node.op_type == 'Add':
        print("Add节点输入:", node.input, "输出:", node.output)

内容的提问来源于stack exchange，提问作者JJJ