Closed clover978 closed 3 years ago
When I check the onnx
file generated with batchsize = 64, I found the last few layers in network head is not right.
The first dimension of reshape layers ought to keep unchanged, but it's not.
Those reshape layers are defined in https://github.com/linghu8812/nanodet/blob/main/nanodet/model/head/nanodet_head.py
if torch.onnx.is_in_onnx_export():
cls_score = torch.sigmoid(cls_score).reshape(1, self.num_classes, -1).permute(0, 2, 1)
bbox_pred = bbox_pred.reshape(1, (self.reg_max+1)*4, -1).permute(0, 2, 1)
bbox_pred = F.softmax(bbox_pred.reshape(1, -1, self.reg_max + 1), dim=2)
bbox_pred = self.linear(bbox_pred).reshape(1, -1, 4)
return cls_score, bbox_pred
https://github.com/linghu8812/nanodet should add batch size when export onnx, will fix this later.
To make batch inference produce right result. We need to pass variable batchsize to reshape fuction. modify https://github.com/linghu8812/nanodet/blob/main/nanodet/model/head/nanodet_head.py#L136 works.
if torch.onnx.is_in_onnx_export():
batchsize = x.data.size(0)
cls_score = torch.sigmoid(cls_score).reshape(batchsize , self.num_classes, -1).permute(0, 2, 1)
bbox_pred = bbox_pred.reshape(1, (self.reg_max+1)*4, -1).permute(0, 2, 1)
bbox_pred = F.softmax(bbox_pred.reshape(batchsize , -1, self.reg_max + 1), dim=2)
bbox_pred = self.linear(bbox_pred).reshape(batchsize , -1, 4)
return cls_score, bbox_pred
fixed by above PR. closed.
When I run nanodet tensorrt inference with batchsize = 1, following authur's instruction here. The result is quiet normal. However, when I just switch to batchsize = 64, and regenerate
onnx
file , modifyconfig
file. The results are in a mess.