目录
前言
自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》,链接。记录下个人学习笔记,仅供自己参考
本次课程我们来学习课程第五章—TensorRT API 的基本使用,一起来学习 TensorRT 模型搭建的模块化
课程大纲可以看下面的思维导图
0. 简述
本小节目标:学习 TensorRT 模型搭建的模块化设计思想
今天我们来讲第五章节第六小节—5.6-build-sub-graph 这个案例,上节课中我们学习了利用 C++ API 搭建网络的各个 layer 比如 conv、bn、reshape 等等,其实我们就可以利用这些 layer 搭建一个完整的模型,例如 resnet、yolo、vit 等等,理论上我们都是可以实现的
但是整个模型的搭建其实还是有非常多的地方需要注意的,我们目前只学习了搭建某个 layer,那这些都是模型的一些小零件,如果按照一个个 layer 去搭建其实是非常麻烦的,因为一个模型可能就有上百层,一层层搭建太过于繁琐,整个代码也非常冗长,所以我们这里学习 pytorch 搭建的方式采用模块化设计的思想,先搭建一个个 module 然后利用这些 module 搭建 network
下面我们开始本次课程的学习🤗
1. 案例运行
在正式开始课程之前,博主先带大家跑通 5.6-build-sub-graph 这个小节的案例🤗
源代码获取地址:https://github.com/kalfazed/tensorrt_starter
首先大家需要把 tensorrt_starter 这个项目给 clone 下来,指令如下:
git clone https://github.com/kalfazed/tensorrt_starter.git
也可手动点击下载,点击右上角的 Code
按键,将代码下载下来。至此整个项目就已经准备好了。也可以点击 here 下载博主准备好的源代码(注意代码下载于 2024/7/14 日,若有改动请参考最新)
整个项目后续需要使用的软件主要有 CUDA、cuDNN、TensorRT、OpenCV,大家可以参考 Ubuntu20.04软件安装大全 进行相应软件的安装,博主这里不再赘述
假设你的项目、环境准备完成,下面我们来一起运行 5.6 小节案例代码
开始之前我们需要创建几个文件夹,在 tensorrt_starter/chapter5-tensorrt-api-basics/5.6-build-sub-graph 小节中韩君老师已经创建了 models 文件夹,并且在 models 文件夹下有提供的 weights 文件夹,我们只需要在 models 文件夹下新建一个 onnx 和 engine 文件夹即可
创建完后 5.6 小节整个目录结构如下:
虽然这里已经提供了各种 weights 文件,但是我们还是执行下对应的 python 脚本,方便对比 python 和 c++ 的结果,先进入到 5.6 小节中:
cd tensorrt_starter/chapter5-tensorrt-api-basics/5.6-build-sub-graph
执行如下指令:
python src/python/export_cbr.py
Note:大家需要准备一个虚拟环境,安装好 torch、onnx、onnxsim 等第三方库
输出如下:
生成好的 onnx 模型文件保存在 models/onnx 文件夹下,对应的 weights 文件保存在 models/weights 文件夹下,大家可以查看
接着我们需要加载 weights 利用自己搭建的 network 生成对应的 engine,在此之前我们需要修改下整体的 Makefile.config,指定一些库的路径:
# tensorrt_starter/config/Makefile.config
# CUDA_VER := 11
CUDA_VER := 11.6
# opencv和TensorRT的安装目录
OPENCV_INSTALL_DIR := /usr/local/include/opencv4
# TENSORRT_INSTALL_DIR := /mnt/packages/TensorRT-8.4.1.5
TENSORRT_INSTALL_DIR := /home/jarvis/lean/TensorRT-8.6.1.6
Note:大家查看自己的 CUDA 是多少版本,修改为对应版本即可,另外 OpenCV 和 TensorRT 修改为你自己安装的路径即可
然后我们还要简单修改下源码,在 src/cpp/main.cpp 中默认使用的 weights 是 sample_c2f.weights,我们修改为 sample_cbr.weights,修改如下所示:
# src/cpp/main.cpp
int main(int argc, char const *argv[])
{
Model model("models/weights/sample_cbr.weights", Model::precision::FP32);
// Model model("models/weights/sample_c2f.weights", Model::precision::FP16);
...
}
接着我们就可以来执行编译,指令如下:
make -j64
输出如下:
接着执行:
./trt-infer
输出如下:
我们这里通过手动构建一个 conv+bn+relu 的 module 并加载相应的 weights 权重完成模块的构建和推理,可以看到和 python 推理结果保持一致
Note:博主这里也准备了其它的模块和相应的权重,大家可以点击 here 下载,然后运行代码看下其它模块的搭建过程
如果大家能够看到上述输出结果,那就说明本小节案例已经跑通,下面我们就来看看具体的代码实现
2. 代码分析
2.1 main.cpp
我们先从 main.cpp 看起:
#include <iostream>
#include <memory>
#include "utils.hpp"
#include "model.hpp"
using namespace std;
int main(int argc, char const *argv[])
{
/*
* 这里面依次举几个例子来进行展示, 对应的输入和输出也会不一样
* sample_cbr: conv + BN + ReLU: input shape: [1x1x5x5], output shape: [1x1x3x3]
* sample_resBlock: ---: input shape: [1x1x5x5], output shape: [1x3x5x5]
* sample_convBNSiLU: conv + BN + SeLU: input shape: [1x1x5x5], output shape: [1x3x5x5]
* sample_c2f: ---: input shape: [1x1x5x5], output shape: [1x4x5x5]
*/
// Model model("models/weights/sample_cbr.weights", Model::precision::FP32);
// Model model("models/weights/sample_cbr.weights", Model::precision::FP16);
// Model model("models/weights/sample_resBlock.weights", Model::precision::FP32);
// Model model("models/weights/sample_resBlock.weights", Model::precision::FP16);
// Model model("models/weights/sample_convBNSiLU.weights", Model::precision::FP32);
// Model model("models/weights/sample_c2f.weights", Model::precision::FP32);
Model model("models/weights/sample_c2f.weights", Model::precision::FP16);
if(!model.build()){
LOGE("fail in building model");
return 0;
}
if(!model.infer()){
LOGE("fail in infering model");
return 0;
}
return 0;
}
Model 的构造函数和上节课差不多,接收一个 weights 此外还额外添加了一个 precision 精度参数,可以设置为 FP32 或者 FP16,我们简单看下其构造函数的实现:
Model::Model(string path, precision prec){
if (getFileType(path) == ".onnx")
mOnnxPath = path;
else if (getFileType(path) == ".weights")
mWtsPath = path;
else
LOGE("ERROR: %s, wrong weight or model type selected. Program terminated", getFileType(path).c_str());
if (prec == precision::FP16) {
mPrecision = nvinfer1::DataType::kHALF;
} else if (prec == precision::INT8) {
mPrecision = nvinfer1::DataType::kINT8;
} else {
mPrecision = nvinfer1::DataType::kFLOAT;
}
mEnginePath = getEnginePath(path, prec);
}
其中 precision 精度是一个枚举类型,根据传入的 prec 参数来指定 mPrecision 精度,其类型是 nvinfer1::DataType
2.2 model.cpp
我们重点来看下 build 接口:
bool Model::build() {
if (mOnnxPath != "") {
return build_from_onnx();
} else {
return build_from_weights();
}
}
和上小节案例一样,也是调用的 build_from_weights
函数来 build model
我们重点来看下该函数的实现:
if (fileExists(mEnginePath)){
LOG("%s has been generated!", mEnginePath.c_str());
return true;
} else {
LOG("%s not found. Building engine...", mEnginePath.c_str());
}
mWts = loadWeights();
// 这里和之前的创建方式是一样的
Logger logger;
auto builder = make_unique<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
auto config = make_unique<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
auto network = make_unique<nvinfer1::INetworkDefinition>(builder->createNetworkV2(1));
前面的这部分和上小节案例一模一样,都是先加载 weights 然后创建 build,创建 config,创建 network
// 根据不同的网络架构创建不同的TensorRT网络,这里使用几个简单的例子
if (mWtsPath == "models/weights/sample_cbr.weights") {
network::build_cbr(*network, mPrecision, mWts);
} else if (mWtsPath == "models/weights/sample_resBlock.weights") {
network::build_resBlock(*network, mPrecision, mWts);
} else if (mWtsPath == "models/weights/sample_convBNSiLU.weights") {
network::build_convBNSiLU(*network, mPrecision, mWts);
} else if (mWtsPath == "models/weights/sample_c2f.weights") {
network::build_C2F(*network, mPrecision, mWts);
} else {
return false;
}
接着根据不同的 module 调用不同的函数,只不过这里通过 network
命名空间来创建的,下面我们来看看 network
这个命名空间所做的工作
2.3 network.hpp
network.hpp 头文件代码如下:
#ifndef __NETWORK_HPP__
#define __NETWORK_HPP__
#include <NvInfer.h>
#include <string>
#include <map>
#include <memory>
#include <model.hpp>
namespace network {
namespace parser {
nvinfer1::IShuffleLayer* addReshape(
std::string layer_name,
nvinfer1::ITensor& input,
std::vector<int> dims,
std::vector<int> perm,
nvinfer1::INetworkDefinition& network);
nvinfer1::IShuffleLayer* addPermute(
std::string layer_name,
nvinfer1::ITensor& input,
std::vector<int> perm,
nvinfer1::INetworkDefinition& network);
nvinfer1::IFullyConnectedLayer* addFullyConnected(
std::string layer_name,
nvinfer1::ITensor& input,
int output_channel,
nvinfer1::INetworkDefinition& network,
std::map<std::string, nvinfer1::Weights> weights);
nvinfer1::IScaleLayer* addBatchNorm(
std::string layer_name,
nvinfer1::ITensor& input,
nvinfer1::INetworkDefinition& network,
std::map<std::string, nvinfer1::Weights> weights);
nvinfer1::IConvolutionLayer* addConv2d(
std::string layer_name,
nvinfer1::ITensor& input,
int kernel_size, int output_channel, int stride, int pad,
nvinfer1::DataType prec,
nvinfer1::INetworkDefinition& network,
std::map<std::string, nvinfer1::Weights> weights);
nvinfer1::IActivationLayer* addActivation(
std::string layer_name,
nvinfer1::ITensor& input,
nvinfer1::ActivationType type,
nvinfer1::INetworkDefinition& network);
nvinfer1::IElementWiseLayer* addElementWise(
std::string layer_name,
nvinfer1::ITensor& input1,
nvinfer1::ITensor& input2,
nvinfer1::ElementWiseOperation type,
nvinfer1::INetworkDefinition& network);
nvinfer1::IElementWiseLayer* addConvBNSiLU(
std::string layer_name,
nvinfer1::ITensor& input,
int kernel_size,
int output_channel,
int stride,
int pad,
nvinfer1::DataType prec,
nvinfer1::INetworkDefinition& network,
std::map<std::string, nvinfer1::Weights> weights);
nvinfer1::ILayer* addC2F(
std::string layer_name,
nvinfer1::ITensor& input,
int output_channel,
nvinfer1::DataType prec,
nvinfer1::INetworkDefinition& network,
std::map<std::string, nvinfer1::Weights> weights);
} // namespace parser
// void build_linear(nvinfer1::INetworkDefinition& network, std::map<std::string, nvinfer1::Weights> mWts);
// void build_conv(nvinfer1::INetworkDefinition& network, std::map<std::string, nvinfer1::Weights> mWts);
// void build_permute(nvinfer1::INetworkDefinition& network, std::map<std::string, nvinfer1::Weights> mWts);
// void build_reshape(nvinfer1::INetworkDefinition& network, std::map<std::string, nvinfer1::Weights> mWts);
// void build_batchNorm(nvinfer1::INetworkDefinition& network, std::map<std::string, nvinfer1::Weights> mWts);
void build_cbr(
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType prec,
std::map<std::string, nvinfer1::Weights> weights) ;
void build_resBlock(
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType prec,
std::map<std::string, nvinfer1::Weights> weights) ;
void build_convBNSiLU(
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType prec,
std::map<std::string, nvinfer1::Weights> weights) ;
void build_C2F(
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType prec,
std::map<std::string, nvinfer1::Weights> weights);
}; // namespace network
#endif //__NETWORK_HPP__
这个头文件定义了 network
命名空间,并包含了一系列函数声明,这些函数用来在 TensorRT 中构建不同的模块,此外这个命名空间下还定义了一个子命名空间 parser
,用于在 TensorRT 网络定义中添加各种 layer 和 module
下面我们就来看不同模块案例的具体实现
3. 案例
3.1 sample_cbr
我们先看 sample_cbr 案例,代码如下:
void build_cbr(
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType prec,
map<string, nvinfer1::Weights> weights)
{
auto input = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});
auto conv = parser::addConv2d("conv", *input, 3, 3, 1, 0, prec, network, weights);
auto bn = parser::addBatchNorm("norm", *conv->getOutput(0), network, weights);
auto leaky = parser::addActivation("leaky", *bn->getOutput(0), nvinfer1::ActivationType::kLEAKY_RELU, network);
leaky->getOutput(0) ->setName("output0");
network.markOutput(*leaky->getOutput(0));
}
首先通过 network 的 addInput 接口创建输入向量,接着调用 parser 命名空间下的模块函数 addConv2d、addBatchNorm、addActivation 分别添加 conv、bn 和 leaky relu,与上个案例相比,conv、bn 和 relu 的具体实现被封装到了 parser 下面,使得代码更加简洁和模块化
下面是 parser 中各个 layer 的具体实现:
nvinfer1::IScaleLayer* addBatchNorm(
string layer_name,
nvinfer1::ITensor& input,
nvinfer1::INetworkDefinition& network,
std::map<std::string, nvinfer1::Weights> weights)
{
// 因为TensorRT内部没有BatchNorm的实现,但是我们只要知道BatchNorm的计算原理,就可以使用IScaleLayer来创建BN的计算
// IScaleLayer主要是用在quantization和dequantization,作为提前了解,我们试着使用IScaleLayer来搭建于一个BN的parser
// IScaleLayer可以实现: y = (x * scale + shift) ^ pow
float* gamma = (float*)weights[layer_name + ".weight"].values;
float* beta = (float*)weights[layer_name + ".bias"].values;
float* mean = (float*)weights[layer_name + ".running_mean"].values;
float* var = (float*)weights[layer_name + ".running_var"].values;
float eps = 1e-5;
int count = weights[layer_name + ".running_var"].count;
float* scales = (float*)malloc(count * sizeof(float));
float* shifts = (float*)malloc(count * sizeof(float));
float* pows = (float*)malloc(count * sizeof(float));
// 这里具体参考一下batch normalization的计算公式,网上有很多
for (int i = 0; i < count; i ++) {
scales[i] = gamma[i] / sqrt(var[i] + eps);
shifts[i] = beta[i] - (mean[i] * gamma[i] / sqrt(var[i] + eps));
pows[i] = 1.0;
}
// 将计算得到的这些值写入到Weight中
auto scales_weights = nvinfer1::Weights{nvinfer1::DataType::kFLOAT, scales, count};
auto shifts_weights = nvinfer1::Weights{nvinfer1::DataType::kFLOAT, shifts, count};
auto pows_weights = nvinfer1::Weights{nvinfer1::DataType::kFLOAT, pows, count};
// 创建IScaleLayer并将这些weights传进去,这里使用channel作为scale model
auto bn = network.addScale(input, nvinfer1::ScaleMode::kCHANNEL, shifts_weights, scales_weights, pows_weights);
bn->setName(layer_name.c_str());
LOGV("%s, %s", bn->getName(), (printDims(bn->getOutput(0)->getDimensions())).c_str());
return bn;
}
nvinfer1::IConvolutionLayer* addConv2d(
string layer_name,
nvinfer1::ITensor& input,
int kernel_size,
int output_channel,
int stride,
int pad,
nvinfer1::DataType prec,
nvinfer1::INetworkDefinition& network,
std::map<std::string, nvinfer1::Weights> weights)
{
auto conv = network.addConvolutionNd(
input, output_channel,
nvinfer1::DimsHW{kernel_size, kernel_size},
weights[layer_name + ".weight"],
weights[layer_name + ".bias"]);
conv->setName(layer_name.c_str());
conv->setStride(nvinfer1::DimsHW(stride, stride));
conv->setPaddingNd(nvinfer1::DimsHW(pad, pad));
// 注意,这里setPrecision需要跟config->setFlag配合使用,否则无效
conv->setPrecision(prec);
LOGV("%s, %s", conv->getName(), (printDims(conv->getOutput(0)->getDimensions())).c_str());
return conv;
}
nvinfer1::IActivationLayer* addActivation(
string layer_name,
nvinfer1::ITensor& input,
nvinfer1::ActivationType type,
nvinfer1::INetworkDefinition& network)
{
auto act = network.addActivation(input, type);
act->setName(layer_name.c_str());
LOGV("%s, %s", act->getName(), (printDims(act->getOutput(0)->getDimensions())).c_str());
return act;
}
这些函数封装了 TensorRT 底层 API 的调用,其实也就是上节课的各个 layer 的案例,这种方式简化了模型构建过程,存在以下优点:
- 模块化设计:通过将各个网络层的添加逻辑封装成函数,提高了代码的可读性和可维护性
- 易于使用:用户只需调用这些函数并提供必要的参数,即可轻松得向 TensorRT 网络定义中添加相应的层
- 扩展性:这种设计模式允许轻松地扩展其他类型的网络层或功能,满足不同的模型需求
该案例执行后的输出如下所示:
FP16 精度的输出如下所示:
对比下 python 结果:
可以看到输出基本都相同,这个就是 sample_cbr 案例
3.2 sample_resBlock
下面我们来看看 sample_resBlock 案例,代码如下:
// 做一个residual block:
// conv0
// / \
// / conv1
// | |
// | bn1
// | |
// | relu1
// | |
// | |
// | conv2
// | |
// | bn2
// \ /
// \ /
// add2
// |
// relu2
//
void build_resBlock(
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType prec,
map<string, nvinfer1::Weights> weights)
{
auto data = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});
auto conv0 = parser::addConv2d("conv0", *data, 3, 3, 1, 1, prec, network, weights);
auto conv1 = parser::addConv2d("conv1", *conv0->getOutput(0), 3, 3, 1, 1, prec, network, weights);
auto bn1 = parser::addBatchNorm("norm1", *conv1->getOutput(0), network, weights);
auto relu1 = parser::addActivation("relu1", *bn1->getOutput(0), nvinfer1::ActivationType::kRELU, network);
auto conv2 = parser::addConv2d("conv2", *relu1->getOutput(0), 3, 3, 1, 1, prec, network, weights);
auto bn2 = parser::addBatchNorm("norm2", *conv2->getOutput(0), network, weights);
auto add2 = parser::addElementWise("add2", *conv0->getOutput(0), *bn2->getOutput(0), nvinfer1::ElementWiseOperation::kSUM, network);
auto relu2 = parser::addActivation("relu2", *add2->getOutput(0), nvinfer1::ActivationType::kRELU, network);
relu2->getOutput(0) ->setName("output0");
network.markOutput(*relu2->getOutput(0));
}
build_resBlock
函数用于在 TensorRT 中构建一个 Residual Block 残差块,我们知道其结构后就可以像搭积木一样把它组合起来就行了,其中 conv、bn、activation 我们上个案例已经分析过了,这里我们分析下 ElementWise 这个 layer,具体实现如下:
nvinfer1::IElementWiseLayer* addElementWise(
string layer_name,
nvinfer1::ITensor& input1,
nvinfer1::ITensor& input2,
nvinfer1::ElementWiseOperation type,
nvinfer1::INetworkDefinition& network)
{
auto ew = network.addElementWise(input1, input2, type);
ew->setName(layer_name.c_str());
LOGV("%s, %s", ew->getName(), (printDims(ew->getOutput(0)->getDimensions())).c_str());
return ew;
}
addElementWise
函数用于在 TensorRT 中添加一个 IElementWiseLayer
层,该层执行元素级操作(element-wise operations),即在逐元素基础上对输入张量进行操作。这种操作对于实现诸如加法、乘法、最小值、最大值等功能非常有用,尤其是在实现残差连接时,其中:
- layer_name:用于标识该层的名称
- input1: 第一个输入张量
- input2: 第二个输入张量
- type: 元素级操作的类型,由
nvinfer1::ElementWiseOperation
枚举类型指定,它包括- kSUM: 逐元素加法
- kPROD: 逐元素乘法
- …
该案例执行后的输出如下所示:
FP16 精度的输出如下所示:
对比下 python 结果:
可以看到输出基本都相同,这个就是 sample_resBlock 案例
3.3 sample_convBNSiLU
下面我们来看看 sample_convBNSiLU 案例,代码如下:
// 做一个conv + bn + SiLU: (yolov8的模块测试)
// conv
// |
// bn
// / \
// | |
// | sigmoid
// \ /
// \ /
// Mul
//
void build_convBNSiLU(
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType prec,
map<string, nvinfer1::Weights> weights)
{
auto data = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});
auto silu = parser::addConvBNSiLU("", *data, 3, 3, 1, 1, prec, network, weights);
silu->getOutput(0) ->setName("output0");
network.markOutput(*silu->getOutput(0));
}
build_convBNSiLU
函数用于在 TensorRT 中构建一个包含卷积(Conv)、批量归一化(BatchNorm)和 SiLU 激活(Swish 激活函数)的模块,这是 YOLOv8 网络结构中的一个常见模块。它主要是通过 parser 中的 addConvBNSiLU
函数来实现的,代码如下:
nvinfer1::IElementWiseLayer* addConvBNSiLU(
string layer_name,
nvinfer1::ITensor& input,
int kernel_size,
int output_channel,
int stride,
int pad,
nvinfer1::DataType prec,
nvinfer1::INetworkDefinition& network,
std::map<std::string, nvinfer1::Weights> weights)
{
auto conv = addConv2d(layer_name + "conv", input, kernel_size, output_channel, stride, pad, prec, network, weights);
auto bn = addBatchNorm(layer_name + "norm", *conv->getOutput(0), network, weights);
auto sigmoid = addActivation(layer_name + "sigmoid", *bn->getOutput(0), nvinfer1::ActivationType::kSIGMOID, network);
auto mul = addElementWise(layer_name + "mul", *bn->getOutput(0), *sigmoid->getOutput(0), nvinfer1::ElementWiseOperation::kPROD, network);
return mul;
}
实现也非常简单,其中的 conv、bn、activation 以及 elementwise 我们前面都已经讲过了,值得注意的是这里的 ElementWiseOperation
类型 kPROD 也就是我们前面提到的逐元素相乘,这是因为 SiLU(SWish) 激活函数的公式是
S
i
L
U
(
x
)
=
x
⋅
σ
(
x
)
SiLU(x) = x \cdot \sigma(x)
SiLU(x)=x⋅σ(x),其中的
σ
(
x
)
\sigma(x)
σ(x) 是 Sigmodi 函数
该案例执行后的输出如下所示:
对比下 python 结果:
可以看到输出基本都相同,这个就是 sample_convBNSiLU 案例
3.4 sample_c2f
下面我们来看看 sample_c2f 案例,代码如下:
// 做一个C2F: (yolov8的模块测试)
// input
// |
// convBNSiLU (n * ch)
// / | \
// / | \
// | | |
// | | convBNSiLU ( 0.5n * ch)
// | | |
// | | convBNSiLU ( 0.5n * ch)
// | | |
// | \ /
// | \ /
// | add (0.5n * ch)
// \ /
// \ /
// Concat (1.5n * ch)
// |
// convBNSiLU (n * ch)
void build_C2F(
nvinfer1::INetworkDefinition& network,
nvinfer1::DataType prec,
map<string, nvinfer1::Weights> weights)
{
auto data = network.addInput("input0", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4{1, 1, 5, 5});
auto c2f = parser::addC2F("", *data, 4, prec, network, weights);
c2f->getOutput(0) ->setName("output0");
network.markOutput(*c2f->getOutput(0));
}
C2f
也是 YOLOv8 网络结构中的一个常见模块,它主要是通过 parser 中的 addC2F
函数来实现的,它的结构如下图所示:
具体实现代码如下:
nvinfer1::ILayer* addC2F(
string layer_name,
nvinfer1::ITensor& input,
int output_channel,
nvinfer1::DataType prec,
nvinfer1::INetworkDefinition& network,
std::map<std::string, nvinfer1::Weights> weights)
{
auto cv1 = addConvBNSiLU(layer_name + "cv1.", input, 1, output_channel, 1, 0, prec, network, weights);
auto dim = cv1->getOutput(0)->getDimensions();
auto slice1 = addSlice(layer_name + "slice1",
*cv1->getOutput(0),
nvinfer1::Dims4{0, 0, 0, 0}, // B, C, H, W (0, 0, 0, 0)
nvinfer1::Dims4{dim.d[0], dim.d[1]/2, dim.d[2], dim.d[3]}, // B, 1/2 * C, H, W
nvinfer1::Dims4{1, 1, 1, 1}, // 1, 1, 1, 1
network);
auto slice2 = addSlice(layer_name + "slice2",
*cv1->getOutput(0),
nvinfer1::Dims4{0, dim.d[1]/2, 0, 0}, // B, C, H, W (0, 1/2 * C, 0, 0)
nvinfer1::Dims4{dim.d[0], dim.d[1]/2, dim.d[2], dim.d[3]}, // B, 1/2 * C, H, W
nvinfer1::Dims4{1, 1, 1, 1}, // 1, 1, 1, 1
network);
auto add = addBottleNeck(layer_name + "m.0.", *slice2->getOutput(0), 2, 2, true, prec, network, weights);
nvinfer1::ITensor* concat2Input[] = {cv1->getOutput(0), add->getOutput(0)};
auto concat2 = addConcat(layer_name + "concat2", concat2Input, 2, network);
auto cv2 = addConvBNSiLU(layer_name + "cv2.", *concat2->getOutput(0), 1, output_channel, 1, 0, prec, network, weights);
return cv2;
}
C2F 模块中包含以下组件:
- ConvBNSiLU
- Slice
- BottleNeck
- Concat
其中的 ConvBNSiLU 我们已经分析过了,我们来看看其它三个组件,先看 slice 切片操作:
nvinfer1::ISliceLayer* addSlice(
string layer_name,
nvinfer1::ITensor& input,
nvinfer1::Dims start,
nvinfer1::Dims size,
nvinfer1::Dims stride,
nvinfer1::INetworkDefinition& network)
{
auto slice = network.addSlice(input, start, size, stride);
slice->setName(layer_name.c_str());
LOGV("%s, %s", slice->getName(), (printDims(slice->getOutput(0)->getDimensions())).c_str());
return slice;
}
addSlice
函数的主要功能是将输入张量按照指定的参数进行切片操作,并返回创建的切片层,具体是通过 network.addSlice
方法完成,其参数解释如下:
- input:输入张量
- start:指定切片的起始位置
- size:指定切片后输出的维度
- stride:指定切片的步长
我们再回过头来看看 C2F 模块中的两个 slice 操作具体做了些什么,代码如下:
auto slice1 = addSlice(layer_name + "slice1",
*cv1->getOutput(0),
nvinfer1::Dims4{0, 0, 0, 0}, // B, C, H, W (0, 0, 0, 0)
nvinfer1::Dims4{dim.d[0], dim.d[1]/2, dim.d[2], dim.d[3]}, // B, 1/2 * C, H, W
nvinfer1::Dims4{1, 1, 1, 1}, // 1, 1, 1, 1
network);
auto slice2 = addSlice(layer_name + "slice2",
*cv1->getOutput(0),
nvinfer1::Dims4{0, dim.d[1]/2, 0, 0}, // B, C, H, W (0, 1/2 * C, 0, 0)
nvinfer1::Dims4{dim.d[0], dim.d[1]/2, dim.d[2], dim.d[3]}, // B, 1/2 * C, H, W
nvinfer1::Dims4{1, 1, 1, 1}, // 1, 1, 1, 1
network);
其中的 slice1 我们后续没用到,我们重点来看 slice2 操作:
- 输入张量:
*cv1->getOutput(0)
,即卷积层cv1
的输出张量。 - 起始位置 (start):
nvinfer1::Dims4{0, dim.d[1]/2, 0, 0}
- 从批次维度的起始位置开始(0)
- 从通道维度的中间位置开始(
dim.d[1]/2
),即从通道数的一半位置开始 - 从高度维度的起始位置开始(0)
- 从宽度维度的起始位置开始(0)
- 切片大小 (size):
nvinfer1::Dims4{dim.d[0], dim.d[1]/2, dim.d[2], dim.d[3]}
dim.d[0]
: 保持批次维度的大小不变dim.d[1]/2
: 通道数的一半,即切出剩下的一半通道dim.d[2]
: 保持高度维度的大小不变dim.d[3]
: 保持宽度维度的大小不变
- 步长 (stride):
nvinfer1::Dims4{1, 1, 1, 1}
- 在每个维度上按步长 1 进行切片操作
slice2 的输出回作为 BottleNeck 的输入,下面我们就一起来看下 BottleNeck 的实现:
// 做一个bottleneck: (yolov8的模块测试)
// input
// / \
// | |
// | convBNSiLU ( 0.5n * ch)
// | |
// | convBNSiLU ( 0.5n * ch)
// | |
// \ /
// \ /
// add (0.5n * ch)
nvinfer1::ILayer* addBottleNeck(
string layer_name,
nvinfer1::ITensor& input,
int ch1, int ch2,
bool shortcut,
nvinfer1::DataType prec,
nvinfer1::INetworkDefinition& network,
std::map<std::string, nvinfer1::Weights> weights)
{
auto silu1 = addConvBNSiLU(layer_name + "cv1.", input, 3, ch1, 1, 1, prec, network, weights);
auto silu2 = addConvBNSiLU(layer_name + "cv2.", *silu1->getOutput(0), 3, ch2, 1, 1, prec, network, weights);
if (shortcut) {
auto add = addElementWise(layer_name + "cv1.add",
input, *silu2->getOutput(0),
nvinfer1::ElementWiseOperation::kSUM, network);
return add;
}
return silu1;
}
该函数包含以下组件:
- ConvBNSiLU
- ElementWise
- 返回值则根据是否启用 shortcut 连接返回不同的输出层
我们来看下 C2F 模块中的 BottleNeck:
auto add = addBottleNeck(layer_name + "m.0.", *slice2->getOutput(0), 2, 2, true, prec, network, weights);
其各个参数含义如下:
- layer_name + “m.0.”:层名称
- *slice2->getOutput(0):
slice2
的输出张量,作为 bottleneck 层的输入 - 2:第一个卷积层的输出通道数
- 2:第二个卷积层的输出通道数
- true:启用 shortcut 连接
- prec:精度
- network:网络定义对象
- weights:权重映射
我们再来看下 Concat 的实现:
nvinfer1::IConcatenationLayer* addConcat(
string layer_name,
nvinfer1::ITensor* input[],
int size,
nvinfer1::INetworkDefinition& network)
{
auto concat = network.addConcatenation(input, size);
concat->setName(layer_name.c_str());
LOGV("%s, %s", concat->getName(), (printDims(concat->getOutput(0)->getDimensions())).c_str());
return concat;
}
该函数在 TensorRT 网络定义中添加一个拼接层,将多个输入张量按通道维度进行拼接,并返回创建的拼接层,具体实现是使用 network.addConcatenation
方法,其参数解释如下:
- input:输入张量组,包含需要拼接的多个张量
- size:输入张量的数量
最后我们再来看下 C2F 模块的 Concat:
nvinfer1::ITensor* concat2Input[] = {cv1->getOutput(0), add->getOutput(0)};
auto concat2 = addConcat(layer_name + "concat2", concat2Input, 2, network);
在 C2F 模块中,Concat 操作用于将 cv1 和 add 的输出张量拼接在一起,也就是第一个 ConvBNSiLU 模块的输出和 BottleNeck 的输出拼接在一起,最后再经过一个 convBNSiLU 就组成了我们的 C2F 模块
该案例执行后的输出如下所示:
FP16 精度的输出如下所示:
对比下 python 结果:
可以看到输出基本都相同,这个就是 sample_c2f 案例
实际上一个网络模型就是这些不同的模块组成的,我们后续在做 INT8 量化时就可以这样逐层分析,看具体哪一些层对精度影响比较严重,逐层或者说逐模块分析,最后把这些对精度影响比较严重的层不让它做量化,让它跑 FP16 的精度就行
总结
本次课程我们主要学习了 TensorRT 模型构建中模块化的思想,这里以 resBlock 以及 YOLOv8 中的几个常见模块为例去讲解了如何搭建一个模块,那通过这些模块我们就可以自己搭建整个网络框架。当然这种模块化的思想也可以用于 INT8 量化时精度下降的原因排查
OK,以上就是 5.6 小节案例的全部内容了,下节我们来学习 5.7 小节自定义插件的构建,敬请期待😄