PyTorch模型转ONNX后再转TensorRT遇到的几个报错

如何将PyTorch模型转为TensorRT

之前有写过详细的文章，介绍了完整的从 PyTorch->ONNX->TensorRT 如何进行转化。此处不再赘述，感兴趣的读者请点击查看：https://editor.csdn.net/md/?articleId=117778308

IndexError: Attribute not found: axes

IndexError: Attribute not found: axes

网上搜了一下Attribute not found这个报错，相关的帖子大多是跟Pad算子有关。后来在CSDN上发现一篇关于axes的报错，感谢这位博主，帮我顺利定位到了问题。帖子的链接也贴在这里：https://blog.csdn.net/ChuiGeDaQiQiu/article/details/119821974

文章里指出这个错误是由于squeeze这个算子有关的，由于squeeze在PyTorch使用过程中，如果不指名squeeze哪个特定轴，就会将tensor中所有dimension为1的轴都进行squeeze操作，恰好满足我在以下定义中需要的功能。

class ClassificationHead(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.initial_block = DownsamplerBlock(128,256)
        
        self.layers = nn.ModuleList()

        for x in range(0, 5):
            self.layers.append(non_bottleneck_1d(256, 0.3, 1))

        self.layers.append(DownsamplerBlock(256, 512))

        for x in range(0, 2):
            self.layers.append(non_bottleneck_1d(512, 0.3,  1))

        self.pool = nn.AdaptiveAvgPool2d((1,1))
        self.fc = nn.Linear(512, num_classes)
        
    def forward(self, input):

        output = self.initial_block(input)
        for layer in self.layers:
            output = layer(output)
        
        output = self.pool(output)
        output = self.fc(output.squeeze())#这一行出错

        return output

这里后来我思考了以下，这边经过一个AdaptiveAvgPool2d后，形成的tensor的dimension应该是[batch_size, channels, 1, 1]，用squeeze的愿意是希望将tensor的最后两个维度去掉，边进行全连接层的计算。

于是将这一行改为：

 output = self.fc(output.flatten(start_dim=1))#这一行出错

flatten的使用可以参考pytorch的官方链接：https://pytorch.org/docs/stable/generated/torch.flatten.html
start_dim=1的参数是限制flatten操作不要从第0个dim开始，而是从第1个dim开始，就能够满足我上面提到的需求。

这一行改变后，报错消失。

MemoryError: std::bad_alloc

在ONNX转TensorRT的过程中，遇到的了第二个报错

MemoryError: std::bad_alloc

trace到最后的报错是

[TensorRT] ERROR: ArgMax_255: at least 2 dimensions are required for input.

这一错误是推测是由于网络中用到了torch.argmax,但tensorRT不支持这一算子？这一点不是非常确定，但去掉了网络中的torch.argmax之后，这个报错就消失了。

ERROR: Try increasing the workspace size with IBuilderConfig::setMaxWorkspaceSize() if using IBuilder::buildEngineWithConfig, or IBuilder::setMaxWorkspaceSize() if using IBuilder::buildCudaEngine

这一错误的完整错误信息是

[TensorRT] ERROR: Try increasing the workspace size with IBuilderConfig::setMaxWorkspaceSize() if using IBuilder::buildEngineWithConfig, or IBuilder::setMaxWorkspaceSize() if using IBuilder::buildCudaEngine.
[TensorRT] ERROR: ../builder/tacticOptimizer.cpp (1820) - TRTInternal Error in computeCosts: 0 (Could not find any implementation for node ConvTranspose_132.)
[TensorRT] ERROR: ../builder/tacticOptimizer.cpp (1820) - TRTInternal Error in computeCosts: 0 (Could not find any implementation for node ConvTranspose_132.)
Completed creating Engine

一开始我主要怀疑问题是由后面的Could not find any implementation for node ConvTranspose_132，网上搜索后，发现问题可能主要出现在第一个错，也就是需要Try increasing the workspace size。

网上的回答主要提醒的是，max_workspace_size的值，将其设置为= 1 << 30。
这边截取部分代码来说明问题：

with trt.Builder(TRT_LOGGER) as builder, \
            builder.create_network(explicit_batch) as network,  \
            trt.OnnxParser(network, TRT_LOGGER) as parser, \
            builder.create_builder_config() as config: # 使用onnx的解析器绑定计算图，后续将通过解析填充计算图
            profile = builder.create_optimization_profile()
            profile.set_shape("inputs", (1,3,512,512),(1,3,512,512),(1,3,512,512))
            config.add_optimization_profile(profile)

            builder.max_workspace_size = 1 << 30  # 预先分配的工作空间大小,即ICudaEngine执行时GPU最大需要的空间
            builder.max_batch_size = max_batch_size # 执行时最大可以使用的batchsize
            builder.fp16_mode = fp16_mode
            builder.int8_mode = int8_mode

注意到这段参数设置中的这一句：

builder.max_workspace_size = 1 << 30  # 预先分配的工作空间大小,即ICudaEngine执行时GPU最大需要的空间

看完这一篇帖子我才注意到其中的问题出在了哪里，帖子链接：https://forums.developer.nvidia.com/t/tensorrt-engine-cannot-be-built-due-to-workspace-size-even-if-its-set-higher/170898/3

I fixed the workspace adjustment to be applied to the config instead of the builder

也就是说，上面这一句要改为

config.max_workspace_size = 1 << 30  # 预先分配的工作空间大小,即ICudaEngine执行时GPU最大需要的空间

才能正确设置GPU最大需要的空间。