PyTorch模型转ONNX后再转TensorRT遇到的几个报错
文章目录
-
- 如何将PyTorch模型转为TensorRT
- IndexError: Attribute not found: axes
- MemoryError: std::bad_alloc
- ERROR: Try increasing the workspace size with IBuilderConfig::setMaxWorkspaceSize() if using IBuilder::buildEngineWithConfig, or IBuilder::setMaxWorkspaceSize() if using IBuilder::buildCudaEngine
如何将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最大需要的空间。