从TensorRT输出维度 动态获取 分类数onnx转trt

netron查看模型，进行区分

 

 

 

// 从TensorRT输出维度 动态获取 分类数
int FormTensorRT::get_num_classes(Dims output_dims) {
// YOLO输出格式: [1, na*(5 + nc), 8400] 或 [1, 8400, na*(5+nc)]
// nc = 分类数量（自动识别！）
//yolov5的格式：[1,25200,7] ,7是5+2， 2是分类的数量，训练是OK与NG两类
int pred_len = 0;
// if (output_dims.d[1] > output_dims.d[2]) {
//pred_len = output_dims.d[1];
// }
// else {
// pred_len = output_dims.d[2];
//}
pred_len = output_dims.d[2];//这里就取7,7=5+2得到2类
int num_classes = pred_len - 5; // 5 = x,y,w,h,obj
return num_classes;
}