TensorRTX 生成Engine笔记 Alexnet

原创

已于 2022-06-22 13:33:13 修改 · 1k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#c++

于 2022-06-22 11:15:57 首次发布

本文介绍了如何利用TensorRT API将AlexNet模型转换为Engine，涉及构建Builder、配置、添加网络层和创建Engine的详细步骤，包括使用CUDA和NvInfer库。重点讲解了APIToModel函数和CreateEngine的实现过程。

A

本文为https://github.com/wang-xinyu/tensorrtx/tree/master/alexnet笔记
当所有plugin都有tensorrt api支持时，不需要用cuda自己写plugin，否则需要

include

NvInfer.h
用于Engine构造和使用所必须要的API
官网提供的c++API里面 Nvinfer.h

cuda_runtime_api.h
用于cudarun必须要的API

logging.h
顾名思义，这是日志组件，用于管理builder, engine 和 runtime 的日志信息。
通常来说，logger 会作为一个必须的参数传递给 builder runtime parser的实例化接口。
该类为 TensorRT 工具和示例提供了一个通用接口来将信息记录到控制台，并支持记录两种类型的消息：

具有相关严重性（信息、警告、错误或内部错误/致命）的调试消息
测试通过/失败消息
与直接发送到 stdout/stderr 相比，让所有样本都使用此类进行日志记录的优势在于，控制样本输出的详细程度和格式的逻辑集中在一个位置。
将来，可以扩展此类以支持将测试结果转储到某种标准格式（例如，JUnit XML）的文件中，并提供额外的元数据（例如，对测试运行的持续时间进行计时）。

Alexnet-python

class alexnet(nn.Module):
    
    def __init__(self, num_classes=1000):
        super(alexnet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(64, 192, kernel_size=5, stride=1, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(192, 384, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True)  ,  
            nn.MaxPool2d(kernel_size=3, stride=2)
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6,6))
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(6*6*256, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes)
        )
    
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = x.view(x.size