其他分享
首页 > 其他分享> > 【模型推理】教你 tensorrt 实现 mish 算子

【模型推理】教你 tensorrt 实现 mish 算子

作者:互联网

欢迎关注我的公众号 [极智视界],获取我的更多笔记分享

O_o>_<o_OO_o~_~o_O

  本文介绍了使用 tensorrt 实现 mish 算子的方法。

  相信做过目标检测的同学对 yolo 肯定比较熟悉了,yolov4是 2020 年初提出的,相继后来有了 yolov5 和其他一些变体,yolov4 中汇集了很多 tricks,其中 mish 激活函数也在其中。mish 在这篇论文《Mish: A Self Regularized Non-Monotonic Activation Function》里进行了详细介绍,这里我做了一些函数本身的介绍以及 tensorrt 怎么样去实现 mish 算子。

文章目录

1、mish 函数数学表达

  mish 的数学表达式如下:
在这里插入图片描述

  函数的图像表达如下,其中:

  让我们来看看 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);
   ...
}

  好了,收工~



扫描下方二维码即可关注我的微信公众号【极智视界】,获取更多AI经验分享,让我们用极致+极客的心态来迎接AI !
在这里插入图片描述

标签:tanh,activation,mish,tensorrt,relu,softplus,算子
来源: https://blog.csdn.net/weixin_42405819/article/details/120312039