如何在Python中获取TorchScript模型的所有层与算子？

一、遍历所有子模块（含无参数层） #

你的原有代码named_parameters()仅返回带可训练参数的模块，要获取所有子模块（包括ReLU、BatchNorm这类无参数层），可以改用named_modules()方法：

import torch

# 加载TorchScript模型
model = torch.jit.load('best5.torchscript')

# 遍历所有子模块
for name, module in model.named_modules():
    # 打印模块名称和类型
    print(f"模块名称: {name}, 模块类型: {type(module).__name__}")

这个方法能列出模型中所有以模块形式存在的组件，但无法捕获函数式算子（比如直接写x + y对应的add操作，或者torch.relu(x)这类函数调用）。

二、解析计算图获取所有算子 #

如果要获取Netron中看到的所有算子（包括函数式操作），需要解析TorchScript模型的计算图：

import torch

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

# 获取模型的计算图
graph = model.graph

# 遍历计算图中的所有节点（算子）
for node in graph.nodes():
    # 算子类型、名称、输入输出信息
    print(f"算子类型: {node.kind()}, 节点名称: {node.name()}")
    # 可选：打印输入输出
    # print(f"输入: {[i.debugName() for i in node.inputs()]}")
    # print(f"输出: {[o.debugName() for o in node.outputs()]}")

关键说明： #

• node.kind()：返回算子的类型标识，比如aten::add、aten::conv2d、prim::Constant等，对应Netron中显示的算子类型。
• node.name()：返回节点的内部名称，可用于区分同一类型的不同算子实例。
• 部分节点（比如prim::Constant）是模型中的常量参数，如果你只想保留计算类算子，可以通过过滤node.kind()来筛选，比如只保留以aten::开头的节点。

三、结合两种方法的完整方案 #

如果需要同时获取模块结构和所有算子，可以将两种方法结合使用：

import torch

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

print("=== 所有子模块 ===")
for name, module in model.named_modules():
    print(f"{name}: {type(module).__name__}")

print("\n=== 所有算子 ===")
for node in model.graph.nodes():
    # 过滤掉常量、占位符等非计算类节点
    if node.kind().startswith('aten::'):
        print(f"{node.kind().split('::')[1]} (节点名: {node.name()})")

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