nanodet batch inference produces messy result

[clover978]

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 , modify config file. The results are in a mess.

[clover978]

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

[linghu8812]

https://github.com/linghu8812/nanodet should add batch size when export onnx, will fix this later.

[clover978]

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

[clover978]

fixed by above PR. closed.