使用TensorRT加速ScaledYOLOv4
很多人都写过TensorRT版本的yolo了,我也来写一个。具体代码可参见我的github
测试环境
ubuntu 18.04
pytorch 1.7.1
jetpack 4.4
CUDA 11.0
TensorRT7.1
快速开始
1、生成onnx模型
git clone --branch yolov4-csp https://github.com/WongKinYiu/ScaledYOLOv4
git clone https://github.com/talebolano/TensorRT-Scaled-YOLOv4
cp TensorRT-Scaled-YOLOv4/script/ScaledYOLOv4-csp/* ScaledYOLOv4/
下载yolov4-csp.weights到ScaledYOLOv4/
cd ScaledYOLOv4
python3 export.py
2、编译
cd ../TensorRT-Scaled-YOLOv4
mkdir build
cd build
cmake ..
make -j8
3、转换onnx模型到trt模型
./makeCudaEngine -i ../../ScaledYOLOv4/yolov4-csp.onnx -o yolov4-csp.trt
4、测试
./inferYoloCuda -e yolov4-csp.trt -i 你的图片 -show -save
速度效果
Mode | GPU | inference time | Ap |
---|---|---|---|
FP16 | V100 | 10ms | - |
FP16 | xavier | 35ms | - |
使用mish插件层
1、下载TensorRT的开源版,在builtin_op_importers.cpp中注册Mish插件。将以下代码粘贴到builtin_op_importers.cpp底部
DEFINE_BUILTIN_OP_IMPORTER(Mish)
{
ASSERT(inputs.at(0).is_tensor(), nvonnxparser::ErrorCode::kUNSUPPORTED_NODE); // input
std::vector<nvinfer1::ITensor*> tensors;
nvinfer1::ITensor* input = &convertToTensor(inputs.at(0),ctx);
tensors.push_back(input);
const std::string pluginName = "Mish_TRT";
const std::string pluginVersion = "001";
std::vector<nvinfer1::PluginField> f;
const auto mPluginRegistry = getPluginRegistry();
const auto pluginCreator
= mPluginRegistry->getPluginCreator(pluginName.c_str(), pluginVersion.c_str(), "");
nvinfer1::PluginFieldCollection fc;
fc.nbFields = f.size();
fc.fields = f.data();
nvinfer1::IPluginV2* plugin = pluginCreator->createPlugin(node.name().c_str(), &fc);
ASSERT(plugin != nullptr && "Mish plugin was not found in the plugin registry!",
ErrorCode::kUNSUPPORTED_NODE);
nvinfer1::IPluginV2Layer* layer = ctx->network()->addPluginV2(tensors.data(), tensors.size(), *plugin);
RETURN_ALL_OUTPUTS(layer);
}
2、在ScaledYOLOv4/models/models.py中使用MishImplementation()替代Mish(),将
modules.add_module('activation', Mish())
替换为
modules.add_module('activation', MishImplementation())
3、生成onnx模型并转换为trt模型
python3 export.py
./makeCudaEngine -i ../../ScaledYOLOv4/yolov4-csp.onnx -o yolov4-csp.trt
4、测试
./inferYoloCuda -e yolov4-csp.trt -i 你的图片 -show -save
注意:在xavier中使用mish插件要比不使用慢10ms以上。