一起养成写作习惯！这是我参与「掘金日新计划 · 4 月更文挑战」的第11天，点击查看活动详情。
欢迎关注我的公众号 [极智视界]，获取我的更多笔记分享
大家好，我是极智视界。本文介绍了使用 tensorrt 实现 mish 算子的方法。
相信做过目标检测的同学对 yolo 肯定比较熟悉了，yolov4是 2020 年初提出的，相继后来有了 yolov5 和其他一些变体，yolov4 中汇集了很多 tricks，其中 mish 激活函数也在其中。mish 在这篇论文《Mish: A Self Regularized Non-Monotonic Activation Function》里进行了详细介绍，这里我做了一些函数本身的介绍以及 tensorrt 怎么样去实现 mish 算子。
1、mish 函数数学表达
mish 的数学表达式如下：
函数的图像表达如下，其中：
- 蓝色曲线为：mish
- 橙色曲线为：ln(1 + e^(x))
让我们来看看 mish 在 yolov4 里的样子：
也可以把 mish 看成由 tanh 和 softplus 组合起来的。让我们来看下，tanh 的数学表达如下：
softplus 的数学表达如下，softplus 可以看作是 relu 的平滑。
以上对比了 mish、tanh、softplus 的数学表达式，你可以很容易就发现 mish 也可以写成这样：
2、mish 对比 relu
relu 因其能克服梯度消失和加快训练收敛，可以说是最常用的激活函数了，relu 是个分段函数，数学表达式如下：
函数图像表达如下：
这里拿 mish 和 relu 作对比，就是有硬碰硬的感觉。拿论文《Mish: A Self Regularized Non-Monotonic Activation Function》的一些实验进行说明。
这是 relu 和 mish 的梯度对比，可以看到 mish 的梯度更加平滑。
精度方面，在 ImageNet-1K 数据集上对 mish、swish、relu、leaky relu 激活函数对网络精度的提升对了对比，数据如下：
以下是在 MS-COCO 目标检测数据集的对比数据：
从实测的精度提升数据来看，mish 具有十分明显的优势。
性能方面，在 pytorch 框架中，对 relu、softplus、mish、mish-cuda (RTX-2070) 在 fp32 和 fp16 精度下进行了性能对比，数据如下，可以看到 relu 在推理效率上要比 mish 快，mish-cuda 在用 cuda 进行优化后性能能提升很多。
3、tensorrt 实现 mish 算子
先让我们来看一下 tensorrt API 直接支持的激活函数算子：
//!
//! \enum ActivationType
//!
//! \brief Enumerates the types of activation to perform in an activation layer.
//!
enum class ActivationType : int32_t
{
kRELU = 0, //!< Rectified linear activation.
kSIGMOID = 1, //!< Sigmoid activation.
kTANH = 2, //!< TanH activation.
kLEAKY_RELU = 3, //!< LeakyRelu activation: x>=0 ? x : alpha * x.
kELU = 4, //!< Elu activation: x>=0 ? x : alpha * (exp(x) - 1).
kSELU = 5, //!< Selu activation: x>0 ? beta * x : beta * (alpha*exp(x) - alpha)
kSOFTSIGN = 6, //!< Softsign activation: x / (1+|x|)
kSOFTPLUS = 7, //!< Parametric softplus activation: alpha*log(exp(beta*x)+1)
kCLIP = 8, //!< Clip activation: max(alpha, min(beta, x))
kHARD_SIGMOID = 9, //!< Hard sigmoid activation: max(0, min(1, alpha*x+beta))
kSCALED_TANH = 10, //!< Scaled tanh activation: alpha*tanh(beta*x)
kTHRESHOLDED_RELU = 11 //!< Thresholded ReLU activation: x>alpha ? x : 0
};
复制代码可以看到像 relu、sigmoid、tanh ... 这些你都不用自己去写，直接调 trt 的 api 就好了。我们这里的 mish 不是直接支持的，所以用 trt 来实现的话基本有两种思路：
(1) 用已有算子组合，mish 的话可以用 tanh 和 softplus 组合起来；
(2) 用 cuda kernel 实现，用 plugin 注册进 trt 使用；
下面进行介绍。
3.1 已有算子组合实现
这个其实很好写，看看 mish 的数学表达：
所以基本思路就是先调一个 softplus，再调一个 tanh，把 softplus 的结果传给 tanh，tanh 的输出就等效于一个 mish 的输出。关键代码如下：
########### softplus ############
# 需要注意，trt 里的 softplus 长这样：alpha*log(exp(beta*x)+1)
activationSP = network->addActivation(*Layers[inputName], nvinfer1::ActivationType::kSOFTPLUS);
# 将 alpha 和 beta 设置为 1
activationSP->setAlpha(1);
activationSP->setBeta(1);
############# tanh ##############
nvinfer1::ITensor *activationSP_Out = activationSP->getOutput(0);
mish = network->addActivation(*activationSP_Out, nvinfer1::ActivationType::kTANH);
复制代码以上就完成了使用 tensorrt 已有算子组合来实现 mish 操作。
3.2 cuda + plugin 实现
将 mish 进行数学等价转换，转换成如下数学表达：
基本思想就是用 cuda 来直接实现，把原来用 tanh 和 softplus 组合需要两个算子变为一个算子。这里不说用 cuda kernel 怎么去实现了，说一下怎么把 .cu 通过 plugin 注册进 tensorrt 把。
首先你需要有个头，类似这样：
/// mish.h
#include<NvInfer.h>
#include<NvInferPlugin.h>
calss MishLayerPlugin : public IPluginExt
{
void mish_infer(...);
}
复制代码然后是 .cu，里面是算子 gpu_infer 的实现，差不多像这样：
/// mish.cu
#include "mish.h"
__global__ void mish(...)
{
...;
}
void MishLayerPlugin::mish_infer(...)
{
mish<<<xx, xx>>>(...);
}
复制代码最后是 .cpp 里通过 plugin 注册进 tensorrt，差不多像这样：
/// tensort-mish.cpp
#include "mish.h"
void addmish_layer(...)
{
nvinfer1::DataType Dtype;
Dtype = nvinfer1::DataType::kFLOAT;
nvinfer1::IPluginExt *mish = new MishLayerPlugin(xxx, Dtype);
nvinfer1::IPluginLayer *mish_layer = m_network->addPluginExt(&Layers[inputName], 1, *mish);
...
}
复制代码好了，收工~ 以上分享了tensorrt实现mish算子的方法，希望我的分享能对你的学习有一点帮助。
【公众号传送】