Yunfeng's Simple Blog

Yunfeng's Simple Blog

马上订阅 Yunfeng's Simple Blog RSS 更新: https://vra.github.io/atom.xml

TransformerEncoder导出onnx问题解决

2025年1月29日 16:38

1. 问题说明

在使用Pytorch的TransformerEncoder时，导出onnx会将时序长度固定，导致没法采用变长输入，例如下面的简单例子复现了这个问题：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
import torch
import torch.nn as nn


class SimpleTransformer(nn.Module):
    def __init__(self, input_dim=512, num_layers=6, nhead=8):
        super().__init__()
        # 创建Transformer编码器层
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=input_dim,
            nhead=nhead,
            dim_feedforward=2048,
            dropout=0.1,
            activation="relu",
            batch_first=True,  # 使用batch_first格式
        )

        # 创建Transformer编码器
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=num_layers
        )

    def forward(self, x):
        # 输入形状: (batch_size, seq_len, input_dim)
        x = self.input_proj(x)
        output = self.transformer_encoder(x)
        return output


# 实例化模型
model = SimpleTransformer(input_dim=512, num_layers=2, nhead=8)
model.eval()  # 设置为评估模式

# 创建示例输入（batch_size=2, seq_len=10, input_dim=512）
dummy_input = torch.randn(2, 10, 512)

# 导出ONNX模型
torch.onnx.export(
    model,
    (dummy_input,),
    "transformer_encoder.onnx",
    do_constant_folding=True,  # 优化常量折叠
    input_names=["input"],  # 输入节点名称
    output_names=["output"],  # 输出节点名称
    dynamo=True,
)

print("ONNX model exported...
剩余内容已隐藏
查看完整文章以阅读更多
查看完整文章