TensorRT：CNN分类器案例

目录

1.CNN分类器

 1将分类器导出为ONNX格式

2.cpp推理

2.1：build_model()

2.2inference:

1.CNN分类器

学习实现一个完整的CNN分类器案例:

• 模型里面没有softmax操作，在这里采用包裹一层加上softmax节点后再导出模型，这样使得后处理得到的直接就是概率值，避免后处理上再做softmax
• 在c++代码中，则充分采用指针偏移的方式，提升cpu上预处理的效率
• 对于bgr->rgb也避免使用cvtColor实现，而是简单的改变赋值时的索引，提升效率

 1将分类器导出为ONNX格式

1.主干网络：将softmax分类直接加在主干网络resnet18后面

class Classifier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        #使用torchvision自带的与训练模型, 更多模型请参考:https://tensorvision.readthedocs.io/en/master/
        self.backbone = torchvision.models.resnet18(pretrained=True)  
        
    def forward(self, x):
        feature     = self.backbone(x)
        probability = torch.softmax(feature, dim=1)
        return probability
        

2.图像预处理:

# 对每个通道进行归一化有助于模型的训练
imagenet_mean = [0.485, 0.456, 0.406]
imagenet_std  = [0.229, 0.224,