pt转onnx初学者的一些请教

[RCLaplace]

作者你好
我正在学习pt转onnx的步骤，最终目的是学会将pt模型转化到trt模型并使用
我阅读了你的项目后有一些问题想请教一下

目前有两个问题
一个是在我自己拿到一个新的大模型后，如何查看初始pt模型中有哪些模块是比较重要的需要我们转换成onnx模型的，
然后是如何找到这些模块在源代码中的位置以及其功能呢？

感谢大佬！

[Aimol-l]

一般pt模型中的主要模块是在他的模型配置文件中会有，或者是在在其模型类中有写(多个模块组合到一个模型类中)。
pt本质上就是一个dict嘛，他的key通常就是模块的名称，你把一个pt模型的所有key属性打印出来一看就明白用了什么模块了。
要使用这些模块的话，你可以把这些模块单独提取出来使用(从一个大pt中分离出许多小pt),同时也要把这些模块的模型类也要一起单独拿出来，他们是一起使用的。
比如说：sam2_configs/sam2_hiera_small.yaml
你会发现其中model有多个子模块(image_encoder,memory_attention,..) ，而这些子模块也有自己的子模块(trunk,neck,..),target 字段就是这些模块的模型结构了。自顶向下的一套结构，这不就是dict么。
如果没有提供模型配置文件，那就去看顶层模型，去看看这个顶层模型的成员变量中有哪些成员继承了torch.nn.Module，这通常就是子模块了。

[RCLaplace]

老哥你好 我在自己尝试编写转换代码 但是在转换image_encoder模块的时候遇到了一些问题
以下是该函数的原始代码：
class ImageEncoder(nn.Module): # 提取图像的特征表示
def init(
self,
trunk: nn.Module,
neck: nn.Module,
scalp: int = 0,
):
super().init()
self.trunk = trunk
self.neck = neck
self.scalp = scalp
assert (
self.trunk.channel_list == self.neck.backbone_channel_list
), f"Channel dims of trunk and neck do not match. Trunk: {self.trunk.channel_list}, neck: {self.neck.backbone_channel_list}"

def forward(self, sample: torch.Tensor):
    # Forward through backbone
    features, pos = self.neck(self.trunk(sample)) # 将输入样本 sample 传递给主干网络 trunk，然后传递给颈部网络 neck，以提取特征和位置编码。
    if self.scalp > 0:
        # Discard the lowest resolution features
        features, pos = features[: -self.scalp], pos[: -self.scalp]

    src = features[-1]
    
    output = {
        "vision_features": src,
        "vision_pos_enc": pos,
        "backbone_fpn": features,
    }
    return src, pos, features

以下是我自己编写的转换代码：
def export_image_encoder(model):
input_img = torch.randn(1, 3,1024, 1024)
input_img = input_img.to('cuda')
# 测试输入格式是否正确
out = model(input_img)
output_names = ['vision_features', 'vision_pos_enc', 'backbone_fpn']
save_path = r'D:\work\2024\10\T-sam2onnx\SAM2Export\SAM2Export\checkpoints\large\iamge_encoder.onnx'
torch.onnx.export(
model,
input_img,
save_path,
export_params=True,
opset_version=17,
do_constant_folding=False,
input_names=["image"],
output_names=output_names,
)
onnx_model = onnx.load(save_path)
onnx.checker.check_model(onnx_model)
print("image_encoder.onnx model is valid!")
onnx模型在转换导出后验证时没有问题，但是我将该onnx模型调用运行输出时发现其输出的数据结构与原本的pt模型输出结构不同，原本的输出为一个字典，且字典中第三个键中是一个包含3个元素的列表，但用onnx模型输出后对应的数据只包含一个元素，这是怎么回事呢

感谢大佬！

[Aimol-l]

因为onnx比pt更严格。 onnx的输出应该是n个张量，而不是字典。你应该在在forward 函数返回的时候确保是返回的张量，而不是dict。另外，对于要导出onnx的网络，减少出现依赖运行时变量。   su_aimol ***@***.***  

…

------------------ 原始邮件 ------------------ 发件人: ***@***.***>; 发送时间: 2024年11月14日(星期四) 上午10:21 收件人: ***@***.***>; 抄送: ***@***.***>; ***@***.***>; 主题: Re: [Aimol-l/SAM2Export] pt转onnx初学者的一些请教 (Issue #1) 老哥你好 我在自己尝试编写转换代码 但是在转换image_encoder模块的时候遇到了一些问题 以下是该函数的原始代码： class ImageEncoder(nn.Module): # 提取图像的特征表示 def init( self, trunk: nn.Module, neck: nn.Module, scalp: int = 0, ): super().init() self.trunk = trunk self.neck = neck self.scalp = scalp assert ( self.trunk.channel_list == self.neck.backbone_channel_list ), f"Channel dims of trunk and neck do not match. Trunk: {self.trunk.channel_list}, neck: {self.neck.backbone_channel_list}" def forward(self, sample: torch.Tensor): # Forward through backbone features, pos = self.neck(self.trunk(sample)) # 将输入样本 sample 传递给主干网络 trunk，然后传递给颈部网络 neck，以提取特征和位置编码。 if self.scalp > 0: # Discard the lowest resolution features features, pos = features[: -self.scalp], pos[: -self.scalp] src = features[-1] output = { "vision_features": src, "vision_pos_enc": pos, "backbone_fpn": features, } return src, pos, features 以下是我自己编写的转换代码： def export_image_encoder(model): input_img = torch.randn(1, 3,1024, 1024) input_img = input_img.to('cuda') # 测试输入格式是否正确 out = model(input_img) output_names = ['vision_features', 'vision_pos_enc', 'backbone_fpn'] save_path = r'D:\work\2024\10\T-sam2onnx\SAM2Export\SAM2Export\checkpoints\large\iamge_encoder.onnx' torch.onnx.export( model, input_img, save_path, export_params=True, opset_version=17, do_constant_folding=False, input_names=["image"], output_names=output_names, ) onnx_model = onnx.load(save_path) onnx.checker.check_model(onnx_model) print("image_encoder.onnx model is valid!") onnx模型在转换导出后验证时没有问题，但是我将该onnx模型调用运行输出时发现其输出的数据结构与原本的pt模型输出结构不同，原本的输出为一个字典，且字典中第三个键中是一个包含3个元素的列表，但用onnx模型输出后对应的数据只包含一个元素，这是怎么回事呢 感谢大佬！ — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you commented.Message ID: ***@***.***>

[RCLaplace]

原来如此！ 所以在你编写的ImageEncoder类中的

    vision_pos_enc = backbone_out["vision_pos_enc"] # 有3个tensor
    backbone_fpn = backbone_out["backbone_fpn"]     # 有3个tensor
    pix_feat = backbone_out["vision_features"] # 有1个tensor

    expanded_backbone_out = {
        "backbone_fpn": backbone_fpn,
        "vision_pos_enc": vision_pos_enc,
    }
    for i, feat in enumerate(expanded_backbone_out["backbone_fpn"]):
        expanded_backbone_out["backbone_fpn"][i] = feat.expand(1, -1, -1, -1)
    for i, pos in enumerate(expanded_backbone_out["vision_pos_enc"]):
        expanded_backbone_out["vision_pos_enc"][i] = pos.expand(1, -1, -1, -1)
    
    (_,current_vision_feats,current_vision_pos_embeds,_) = self.prepare_backbone_features(expanded_backbone_out)

    current_vision_feat = current_vision_feats[-1] + self.no_mem_embed
    current_vision_feat2 = current_vision_feat.reshape(64,64,1,256).permute(2, 3, 0, 1) # [1,256,64,64]
    
    # flatten HWxNxC -> NxCxHxW
    high_res_features_0 = current_vision_feats[0].reshape(256,256, 1, 32).permute(2, 3, 0, 1) # [1, 32, 256, 256]
    high_res_features_1 = current_vision_feats[1].reshape(128,128, 1, 64).permute(2, 3, 0, 1) # [1, 64, 128, 128]

    # pix_feat              [1, 256, 64, 64]
    # high_res_features_0   [1, 32, 256, 256]
    # high_res_features_1   [1, 64, 128, 128]
    # current_vision_feat   [1, 256, 64, 64]
    # current_vision_pos_embed2 [4096, 1, 256]
    return pix_feat,high_res_features_0,high_res_features_1,current_vision_feat2,current_vision_pos_embeds[-1]

这部分其实是为了展平数据好让其分变量输出然后再拼接起来对吧