CSDN搬家失败,手动导出markdown后再导入博客园
1、版本不匹配
[E] [TRT] Layer:Where_51's output can not be used as shape tensor.
[E] [TRT] Network validation failed.
[E] Engine creation failed.
[E] Engine set up failed.
这实际是由于 pytorch 与 TensorRT 版本不匹配,我的 TensorRT 是 7.0,pytorch 应该是 1.4,但我用了 1.7
因此需要用 1.7 重新读取权重文件,然后用老的方式保存,再用 onnx 导出
def main():
input_shape = (3, 416, 416)
model_onnx_path = "yolov4tiny.onnx"
# model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
# in_channels=3, out_channels=1, init_features=32, pretrained=True)
model = YoloBody(3, 12).cuda()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dummy_input = torch.randn(1, 3, 416, 416, device=device)
# 用1.7版本读取权重
state_dict = torch.load('logs/Epoch120-Total_Loss0.5324-Val_Loss0.8735.pth', map_location=device)
model.load_state_dict(state_dict)
# 保存成1.4版本支持的格式
torch.save(model.state_dict(), 'logs/for_onnx.pth', _use_new_zipfile_serialization=False)
# Python解释器换成torch1.4的环境,重新读取,导出pytorch1.4版本对应的onnx
state_dict = torch.load('logs/for_onnx.pth', map_location=device)
model.load_state_dict(state_dict)
model.train(False)
inputs = ['input_1']
outputs = ['output_1', 'output_2']
dynamic_axes = {'input_1': {0: 'batch'}, 'output_1': {0: 'batch'}}
torch.onnx.export(model,
dummy_input,
model_onnx_path,
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=inputs, output_names=outputs,
dynamic_axes=None)
这样操作产生的 onnx 文件才能被 TensorRT 转换
使用 TensorRT 的 OSS 工具:
class Upsample(nn.Module):
def __init__(self):
super(Upsample, self).__init__()
def forward(self, x, target_size, inference=False):
assert (x.data.dim() == 4)
# _, _, tH, tW = target_size
if inference:
#B = x.data.size(0)
#C = x.data.size(1)
#H = x.data.size(2)
#W = x.data.size(3)
return x.view(x.size(0), x.size(1), x.size(2), 1, x.size(3), 1).\
expand(x.size(0), x.size(1), x.size(2), target_size[2] // x.size(2), x.size(3), target_size[3] // x.size(3)).\
contiguous().view(x.size(0), x.size(1), target_size[2], target_size[3])
else:
return F.interpolate(x, size=(target_size[2], target_size[3]), mode='nearest')
建议打开 --verbose,转换过程会很慢,verbose 打印日志看着能安心点,要不然盯着屏幕会以为卡死了,发慌
2、上采样 scale 问题
[5] Assertion failed: ctx->tensors().count(inputName)
![[output/attachments/d697b05ccbc0cd8ef1baff4f153b9024_MD5.png]]
YOLO 的上采样阶段,Pytorch 使用 opset=11 的 onnx 会导致 upsample 层里增加一个 constant 节点,所以 TensorFlowRT 转换失败,期间参考 pytorch 经 onnx 使用 TensorRT 部署转换踩坑记录中提到的方法,无效
![[output/attachments/95e6ae1c7b9192edb8b6674198bb354e_MD5.png]]
尝试多个版本的 Pytorch 与 onnx 后,upsample 层的问题依然解决不了,最后参考 https://github.com/Tianxiaomo/pytorch-YOLOv4 这个实现,在 inference 时不使用 torch 自己的插值函数,而是自己重写,成功导出 TensorRT
class Upsample(nn.Module):
def __init__(self):
super(Upsample, self).__init__()
def forward(self, x, target_size, inference=False):
assert (x.data.dim() == 4)
# _, _, tH, tW = target_size
if inference:
#B = x.data.size(0)
#C = x.data.size(1)
#H = x.data.size(2)
#W = x.data.size(3)
return x.view(x.size(0), x.size(1), x.size(2), 1, x.size(3), 1).\
expand(x.size(0), x.size(1), x.size(2), target_size[2] // x.size(2), x.size(3), target_size[3] // x.size(3)).\
contiguous().view(x.size(0), x.size(1), target_size[2], target_size[3])
else:
return F.interpolate(x, size=(target_size[2], target_size[3]), mode='nearest')
3、数据类型错误
Unsupported ONNX data type: DOUBLE (2)
由于上面的改动中,出现了除法,导致 TRT 识别双精度,下图中出现的 cast 节点即为问题。很神奇的是用上面原始的代码不会出现这个问题,换了我自己的模型就有问题了
![[output/attachments/20922cff4c9b9fce378baf949d42b34f_MD5.png]]
根据 https://github.com/onnx/onnx-tensorrt/issues/400#issuecomment-730240546 的说法
![[output/attachments/960f5b584c8b980e4c9713b43a46a8a4_MD5.png]]
尝试了下并没什么用。
为了解决这个问题,直接把 target_size[3] // x.size(3) 换成结果 2 即可,成功。
标签:engine,set,target,onnx,torch,up,output,data,size From: https://www.cnblogs.com/algorithmSpace/p/18200236