<Good Luck To You!> 's work for <TRT-Hackathon-2022-final>
本次复赛我们选择使用TensorRT优化部署的模型是MobileVit，该项工作由Apple的研究者发表在ICLR2022上，项目开源在地址。
- 编译运行代码需要有能正常运行tensorRT的docker（可以从NVIDIA的NGC自行下载），同时运行脚本时必须是root身份。报告中的测试结果是在A10显卡上使用tensorRT8.4GA进行测试的
- 运行
cd build_env && build.sh文件进行代码所需的环境配置 - 运行
cd src/classification && sh autobuild.sh进行mobileVit分类模型的导出及精度与速度测试，直接运行autobuild.sh代表使用fp32且仅进行速度测试，如果需要使用fp16且进行精度测试，请使用autobuild.sh -a 1 -t 1，具体参数含义请见autobuild.sh -h - 运行
cd src/detection && autobuild.sh进行ssd-mobileVit检测模型的导出及精度与速度测试（注意按照sh文件内容修改相应dataset路径），直接运行autobuild.sh代表使用fp32且仅进行速度测试，示例图片测试结果请见src/detection目录下的*result*文件夹，如需运行其他模式请见autobuild.sh -h
说明：如果需要进行分类与检测模型的精度测试，请自行下载ImageNet数据集以及MS-COCO数据集，同时修改
src/ml-cvnets/pretained_models/*.yaml中dataset路径。ml-cvnets/gen_npz.py为生成模拟数据的脚本，用于测试速度，可在分类与检测文件夹中的autobuild.sh查看使用方法
MobileVit结合了CNN和ViT的优势构建了一个轻量级、低延时和移动设备友好的通用基础网络模型。轻量卷积神经网络(CNN)是移动视觉任务的实际应用。CNN的空间归纳偏差允许他们在不同的视觉任务中以较少的参数学习表征，然而CNN在空间上是却是局部的；为了学习全局表征，MobileVit引入了基于自注意力的轻量化的Vision Transformer(ViTs)模块，轻量化主要体现在使用了更小的注意力模块参数L与d。MobileVit能够作为检测、分割等视觉任务的基础网络应用于移动设备之中，均易于部署且能收获到不错的精度。
模型的整体结构如下图所示,其中MV2代表MobileNetV2 block：
MobileVit共有三种不同的网络规模(XXS，XS，S)，参数量逐渐升高，分别为1.3M，2.3M，5.6M，具体如下图所示:
在MS-COCO数据集与ImageNet数据集上对比各个轻量化网络模型的测试结果如下图所示:
当前针对于MobileVit-S网络模型的优化难点总结如下：
- 对于动态batch size的支持。Apple提供的pt与config文件无法直接生成动态batch size的TensorRT engine，需要手动调整源代码中的网络结构以及对算子图进行优化。
- attention、layernorm等模块TensorRT plugin的实现，数据格式支持FP32、TF32以及FP16；使用Nsys分析实现结果，使用FasterTransformer实现MobileVit网络中的Transformer block插件。
- int8以及int8相关TensorRT plugin的实现。轻量化模型在移动端使用int8数据格式能获得更好的加速效果。
- 以MobileVit为基础网络的检测模型相应优化。
修改官网提供的代码导出onnx模型，这个步骤主要是参照官方提供的main_eval.py进行实现，通过将multi_head_attention.py与mobilevit_block.py中reshape操作的-1放在第一维度即可支持动态shape的输入。具体原则如下
-
对于任何用到shape、size返回值的参数，例如：tensor.view(tensor.size(0), -1)，B,C,H,W = x.shape 这类操作，避免直接使用tensor.size的返回值，而是加上int转换，tensor.view(int(tensor.size(0)), -1), B,C,H,W = map(int, x.shape)，断开跟踪。
-
对于reshape、view操作时，-1的指定请放到batch维度。其他维度计算出来即可。batch维度禁止指定为大于-1的明确数字。如果是一维，那么直接指定为-1就好。
在main_export.py中添加生成engine的代码，此处遇到如下问题：
- 在设置网络精度为FP16的情况下，如何在不破坏tensorRT融合结果的情况下单独设置某层的精度为FP32?
设置stict type并设置layer上的精度可以把它设置成fp16或fp32，可参考这个例子（设置成fp16）
- 使用polygraphy找到第一个精度不符合要求的层后，如何继续去寻找下一个精度不符合要求的层？直接切割子图进行分块输出还是使用对修改第一个不符合要求层后生成的engine进行继续调试，命令诸如
polygraphy run decoder.plan --model-type engine --trt --load-inputs=custom_decoder_inputs.json --trt-outputs "788" --fp16 --save-results=_fp16.json
我们根据导师的建议所采用的方法是：在网络层数不是非常深的情况下，逐个单改Transformer网络层中的某一层从fp16到fp32，看精度提高幅度，取提高幅度最高的，作为改变的第一层；然后迭代这一过程，直到精度合格。
使用Nsys对当前模型做性能分析（此时均是FP32结果），结果如下所示:
网络推理一次耗时在110.223ms以上，性能热点主要是Myelin所融合的3个Transformer block部分，分别耗时102.17ms/6.416ms/197.245us。
随后编写LayerNorm、attention插件替换ONNX图中节点，测试后发现速度变慢，使用Nsys进行分析。 以Layernorm结果为例子，如下所示：
网络推理一次耗时在116.492ms以上，并且从结果可以看出我们所替换的插件破坏了Myelin对于Transformer block部分的融合，多了许多小块的ForeignNode节点，表明此插件未能超越trt8.4所作的融合。后续我们还针对nsys的分析结果实现了fastertransformerplugin、swishplugin以及int8 PTQ等，具体过程请见经验与体会
当前模型的的实现步骤主要分为以下两个过程：
- 官方提供的源码中并不支持动态shape的导出过程，我们在
ml-cvnets/cvnets/models/detection/ssd.py添加了export_onnx_forward方法，主要目的是简化onnx模型的输出部分，这样就可以直接在CUDA端操作一个矩阵即可；此外在ml-cvnets/cvnets/modules/ssd.py中的forward方法中做了支持动态shape的相关修改，具体规则与3.1 ONNX模型导出章节的说明一致
# ml-cvnets/cvnets/models/detection/ssd.py 添加如下代码，这里主要是将输出张量进行concat操作，以简化onnx图，output形状为[-1, 85]
def export_onnx_forward(self, x: Tensor, *args, **kwargs):
with torch.no_grad():
confidences, locations, anchors = self.ssd_forward(x, is_prediction=True)
scores = F.softmax(confidences, dim=-1)
boxes = box_utils.convert_locations_to_boxes(
pred_locations=locations,
anchor_boxes=anchors,
center_variance=getattr(self.opts, "model.detection.ssd.center_variance", 0.1),
size_variance=getattr(self.opts, "model.detection.ssd.size_variance", 0.2)
)
return torch.cat([boxes, scores], dim=-1)
# ml-cvnets/cvnets/modules/ssd.py 代码作如下修改，主要是为了切断pytorchshape相关函数的跟踪,并仅将batch_size设置为变化维度
def forward(self, x: Tensor):
batch_size = int(x.shape[0])
if self.proj_layer is not None:
x = self.proj_layer(x)
# [B x C x H x W] --> [B x Anchors * (coordinates + classes) x H x W]
x = self.loc_cls_layer(x)
# [B x Anchors * (coordinates + classes) x H x W] --> [B x H x W x Anchors * (coordinates + classes)]
x = x.permute(0, 2, 3, 1)
# [B x H x W x Anchors * (coordinates + classes)] --> [B x H*W*Anchors X (coordinates + classes)]
mid_dim = int(x.numel() // (self.n_coordinates + self.n_classes) // batch_size)
x = x.contiguous().view(-1, mid_dim, self.n_coordinates + self.n_classes)
# [B x H*W*Anchors X (coordinates + classes)] --> [B x H*W*Anchors X coordinates], [B x H*W*Anchors X classes]
box_locations, box_classes = torch.split(x, [self.n_coordinates, self.n_classes], dim=-1)
return box_locations, box_classes- 编写
main_export_dection.py一键导出模型并生成trt engine。
端到端的推理过程的实现，从数据的预处理和后处理都在GPU上完成。
- 前处理过程 将图片的前处理过程，包含resize、双线性插值与归一化写在一个cuda kernel中，避免多次数据传输与global memroy的访问，部分代码如下所示：
// 双线性插值及resize过程
uint8_t const_value[] = {const_value_st, const_value_st, const_value_st};
float ly = src_y - y_low;
float lx = src_x - x_low;
float hy = 1 - ly;
float hx = 1 - lx;
float w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
float* pdst = dst + dy * dst_width + dx * 3;
uint8_t* v1 = const_value;
uint8_t* v2 = const_value;
uint8_t* v3 = const_value;
uint8_t* v4 = const_value;
if(y_low >= 0){
if (x_low >= 0)
v1 = src + y_low * src_line_size + x_low * 3;
if (x_high < src_width)
v2 = src + y_low * src_line_size + x_high * 3;
}
if(y_high < src_height){
if (x_low >= 0)
v3 = src + y_high * src_line_size + x_low * 3;
if (x_high < src_width)
v4 = src + y_high * src_line_size + x_high * 3;
}
c0 = w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0];
c1 = w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1];
c2 = w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2];
// 归一化过程
c0 = (c0 * norm.alpha - norm.mean[0]) / norm.std[0];
c1 = (c1 * norm.alpha - norm.mean[1]) / norm.std[1];
c2 = (c2 * norm.alpha - norm.mean[2]) / norm.std[2];
- 推理过程 将输入、原始宽高、预处理后的输入及预测框坐标封装为一个结构体
struct Job
{
cv::Mat input;
int ori_width;
int ori_height;
BoxArray output;
shared_ptr<Tensor> mono_tensor;
};随后将数据按照最大batch_size组织并传入反序列化后的trt engine进行推理。
- 后处理过程 后处理过程需要把预测框恢复到原图大小，在做NMS，这两个步骤可以使用cuda实现，核心部分如下：
// 恢复预测框
for(int i = 0; i < num_classes; ++i, ++class_confidence){
if(*class_confidence > confidence_threshold) {
confidence = *class_confidence;
label = i;
int index = atomicAdd(parray, 1);
if(index >= max_objects)
return;
float* pout_item = parray + 1 + index * NUM_BOX_ELEMENT;
*pout_item++ = left;
*pout_item++ = top;
*pout_item++ = right;
*pout_item++ = bottom;
*pout_item++ = confidence;
*pout_item++ = label;
*pout_item++ = 1;
}
}
// nms
float* pcurrent = bboxes + 1 + position * NUM_BOX_ELEMENT;
for(int i = 0; i < count; ++i){
float* pitem = bboxes + 1 + i * NUM_BOX_ELEMENT;
if(i == position || pcurrent[5] != pitem[5]) continue;
if(pitem[4] >= pcurrent[4]){
if(pitem[4] == pcurrent[4] && i < position)
continue;
float iou = box_iou(
pcurrent[0], pcurrent[1], pcurrent[2], pcurrent[3],
pitem[0], pitem[1], pitem[2], pitem[3]
);
if(iou > threshold){
pcurrent[6] = 0; // 1=keep, 0=ignore
return;
}
}
}当前测试环境为A10+Trt8.4。
MobileVit使用imagenet数据集进行精度测试，使用模拟数据测试速度，使用batch_size=128：
| implement | top1-Acc | latancy(ms) |
|---|---|---|
| pytorch-GPU | 78.4 | 155.54 |
| trt-FP32 | 78.28 | 114.52 |
| trt-FP16 | 78.26 | 41.79 |
SSD-MobileVit使用coco数据集进行测试，batch_size=32，每一帧处理时间包含前处理与后处理：
| implement | mAP | FPS |
|---|---|---|
| pytorch | 27.7 | 42.81 |
| trt-FP32 | 26.4 | 281.95 |
| trt-FP16 | 26.4 | 362.17 |
注：尝试了int8 ptq量化方式，推理延时略低，但误差也变大了，故没有将其列入加速效果中。使用trt推理精度损失的主要原因是在检测模型的后处理过程使用了fastnms，在牺牲部分精度的情况下提升了推理速度。
根据nsys分析结果尝试编写一些插件替换性能不足的算子，主要包括几个plugin的编写与int8 PTQ的尝试，plugin部分代码在trt_plugin文件夹中，int8量化部分（需要自行下载imagenet校准数据集）具体代码位于src/Classification/int8_process.py。使用fastertransformer、layernorm、attention plugin后并没有提升速度，使用swish plguin、swish+conv plugin 虽然对精度提升了一个数量级，但推理延时却长了1.5%（bs=128：113.574ms vs 115.275ms），int8 ptq量化虽然推理延时小于fp16结果（bs=128：41.793ms vs 40.303ms）但却带来了比较大的误差。上述各种优化方案的具体测试结果见result_record文件夹，测试环境是Nvidia A10 + trt8.4。
下面简单介绍我们的优化过程，以及在该过程中遇到的一些问题。
通过上一部分分析发现TensorRT8.4已经对模型中Transformer block部分进行了比较高效的融合，如果想超过该结果需要进行更深层次的优化，所以决定开始编写基于FasterTransformer的插件，在编写插件过程中遇到了如下问题
TensorRT C++ api中PluginField的定义如下
class PluginField
{
public:
AsciiChar const* name;
void const* data;
PluginFieldType type;
int32_t length;
PluginField(AsciiChar const* const name_ = nullptr, void const* const data_ = nullptr,
PluginFieldType const type_ = PluginFieldType::kUNKNOWN, int32_t const length_ = 0) noexcept
: name(name_)
, data(data_)
, type(type_)
, length(length_)
{
}
};
通过上面我们可以看到，PluginField中的指针变量采用的是浅拷贝的方式，因此如果想正常使用就需要保证指针指向的变量的生命周期大于等于使用周期。
如下面代码，如果将data_npz定义为循环里面的局部变量，那么循环结束后data_npz占用内存就会被释放，从而导致fieldCollections中新增的PluginField的指针指向一块未被申请的内存，导致结果不正确。
std::vector<PluginField> fieldCollections;
std::string npz_name("npz_path");
cnpy::npz_t data_npz;
int npz_path_len = 0;
std::string npz_path_len_name("npz_path_len");
for(int i = 0; i < fc->nbFields; i++) {
if(npz_path_len_name.compare(fc->fields[i].name) == 0) {
npz_path_len = *((const int*)fc->fields[i].data);
}
}
for(int i = 0; i < fc->nbFields; i++) {
if(npz_name.compare(fc->fields[i].name) == 0) {
string npz_path((const char*)fc->fields[i].data, npz_path_len);
// cnpy::npz_t data_npz = cnpy::npz_load(npz_path);
data_npz = cnpy::npz_load(npz_path);
for(auto& iter: data_npz) {
auto& item = iter.second;
if(0 == type_id) {
fieldCollections.emplace_back(PluginField(iter.first.c_str(), item.data<float>(), PluginFieldType::kFLOAT32, item.num_vals));
} else if(1 == type_id) {
fieldCollections.emplace_back(PluginField(iter.first.c_str(), item.data<half>(), PluginFieldType::kFLOAT16, item.num_vals));
}
}
} else {
fieldCollections.emplace_back(fc->fields[i]);
}
}
PluginField在解析时候是需要指定长度的，但是对于一些可变长度数据来说这个事先是无法确定的。比如onnx中的字符串属性，在传入到TensorRT时，有没有'\0'结尾是不能保证的，因此需要设置一个属性来记录字符串的长度，这也就是上面npz_path_len的用处，对于其他比如float数组也是相同处理方式。
在mobilevit插件编写中，transformer的权重传入有三种方式
- 当作输入传进去
- 作为属性传入，由于属性的数量会随着encoder的层数变化，所以这里需要采用最大属性数量。此外还需要每个传入的数组数据额外添加一个长度属性。
- 把权重npz路径当作属性传入，插件里解析npz来填充权重值
综合考虑源码的改动程度，我们采用第三种方案。
在debug模式下，plugin可以编译成功，但是在构建engine的时候会报LLVM: out of memory的错误。
在release模式下，plugin可以正常编译且正常构建engine。
编写好插件并进行测试后，使用Nsys进行性能分析，发现如下问题：
主要编写了ImageNetEntropyCalibrator，使用imagenet的validation数据进行校准。
在使用了nsys分析fp16 baseline的结果时，发现conv算子(计算密集型)与trt合并的swish算子(访存密集型)所花费的时间差不多，说明trt融合后的swish算子可能出现了memory bound，所以具有一定的优化空间。当前一共实现了两个版本的swish plugin以及swish+conv plugin，分别是基于oneflow-elementwise模板与cudnn API/cudnn_fronted API。