Jetson Nano搭建人脸检测系统: （四）后处理优化

1、目标检测的输出是什么？

  在前面两篇文章中我们使用的人脸检测算法，在经过神经网络模型输出后还进行了一系列的后处理操作，那么这些后处理操作的意义是什么？
  要解决这个问题，我们要先弄懂目标检测算法的网络结构，该模型有两部分输出：类别概率分布和边界框回归。模型的损失函数只是将边界框的回归损失与分类的交叉熵损失相加，通常使用均方误差（MSE）：

  第一个概率分类问题就是图像分类经典问题不再赘述，第二个就是为什么神经网络模型可以做边框回归的问题。我认为主要原因是网络输出不是之前的全连接层，而是一个个特征图，保留了大量的位置信息，可参见这篇论文deep neural networks for object detection(这是最早证明神经网络可以用来回归目标坐标的一篇论文，这篇论文通过实验效果证明了neural networks learned features which are not only good for classification, but also capture strong geometric information)。那么模型的边框回归到底输出的是什么？我们采用4个数来表示一个边界框的坐标，[center x, center y, width, height] ，分别表示边界框的中心坐标以及宽高，但是后面的算法学习的其实是proposal box到gt box的变换系数。
（关于边框回归的发展历史参考One Stage目标检测算法可以直接回归目标的坐标是什么原理？ - 叶不知的回答 - 知乎和深入理解one-stage目标检测算法 - 小小将的文章 - 知乎)
  下面我们以SSD模型的坐标回归为例（与前面的人脸检测算法一致的）来做说明，
先验框位置用表示，其对应的边界框用表示，那么边界框的预测值其实是相对于的转换值：

习惯上，我们称上面这个过程为边界框的编码（encode），预测时，你需要反向这个过程，即进行解码（decode），从预测值中得到边界框的真实位置:

SSD的先验框坐标是归一化到[0,1]，这样它们独立于网格大小（SSD之所以这样是采用了不同大小的网格）,因此得到的坐标是相对于图片的归一化结果，还需乘以图片尺寸，为了得到更精确的坐标还需进行NMS优化。对坐标进行log/exp变换，是为了防止出现负数。SSD的Caffe源码实现中还有trick，那就是设置variance超参数来调整检测值，设置超参数variance，用来对的的4个值进行放缩，此时边界框需要这样解码：

2、TensorRT的常用函数用法实例

TensorRT中已经为我们设计好了常用的Layers，可以通过https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/infer/Graph/Network.html查看。为了能够完成上面的公式，这里介绍几个：

• add_input
• add_scale
• add_slice
• add_constant
• add_elementwise
• add_unary
• add_plugin_v2

2.1 add_input

add_input(self: tensorrt.tensorrt.INetworkDefinition, 
          name: str, 
         dtype: tensorrt.tensorrt.DataType, 
         shape: tensorrt.tensorrt.Dims) → tensorrt.tensorrt.ITensor

功能：为网络添加一个输入层
Parameters :  name  - 层的名字
              dtype - tensor的数据类型，如trt.float32
              shape - tensor的形状，必须小于2^30个元素
Returns:  一个新的tensor

2.2 add_scale

add_scale(self: tensorrt.tensorrt.INetworkDefinition, 
         input: tensorrt.tensorrt.ITensor, 
          mode: tensorrt.tensorrt.ScaleMode, 
         shift: tensorrt.tensorrt.Weights , 
         scale: tensorrt.tensorrt.Weights , 
         power: tensorrt.tensorrt.Weights) → tensorrt.tensorrt.IScaleLayer

功能：控制每个元素缩放大小，其计算公式为
             output → (input*scale+shift)^power

Parameters :  input - 输入tensor，最少有三个维度
               mode - 缩放的模式，如trt.ScaleMode.UNIFORM，表示作用于每一个元素
              shift - Weights变量，公式中的shift值
              scale - Weights变量，公式中的scale值
              power - Weights变量，公式中的power值

如果Weights变量可以得到，那么Weights变量的shape与mode模式相关：
        UNIFORM：形状等于1
        CHANNEL：形状为通道的维度
        ELEMENTWISE：形状与input的形状相同

Returns:  一个新的layer或None

2.3 add_slice

add_slice(self: tensorrt.tensorrt.INetworkDefinition, 
         input: tensorrt.tensorrt.ITensor, 
         start: tensorrt.tensorrt.Dims, 
         shape: tensorrt.tensorrt.Dims, 
        stride: tensorrt.tensorrt.Dims) → tensorrt.tensorrt.ISliceLayer

功能：tensor切片
Parameters :  input - 输入tensor
              start - 起始index
              shape - 输出shape
             stride - 切片步长

Returns:  一个新的layer或None

2.4 add_constant

add_constant(self: tensorrt.tensorrt.INetworkDefinition, 
            shape: tensorrt.tensorrt.Dims, 
          weights: tensorrt.tensorrt.Weights) → tensorrt.tensorrt.IConstantLayer

功能：添加一个常数层，可以把weight对象转变为layer进而变为tensor
Parameters :  shape - 形状
              weights - weight对象

Returns:  一个新的layer或None

2.5 add_elementwise

add_elementwise(self: tensorrt.tensorrt.INetworkDefinition, 
              input1: tensorrt.tensorrt.ITensor, 
              input2: tensorrt.tensorrt.ITensor, 
              op: tensorrt.tensorrt.ElementWiseOperation) → tensorrt.tensorrt.IElementWiseLayer

功能：二元操作
Parameters :  input1（input2） - 输入tensor，形状必须相等
              op - 二元操作符，在ElementWiseOperation中，如：
                    trt.ElementWiseOperation.PROD（乘积）
                    trt.ElementWiseOperation.SUM（加法）

Returns:  一个新的layer或None

2.6 add_unary

add_unary(self: tensorrt.tensorrt.INetworkDefinition,
         input: tensorrt.tensorrt.ITensor, 
         op: tensorrt.tensorrt.UnaryOperation) → tensorrt.tensorrt.IUnaryLayer
功能：一元操作
Parameters :  input1 - 输入tensor，
              op - 一元操作符，在UnaryOperation中，如：
                    trt.UnaryOperation.EXP（自然指数）
                    trt.UnaryOperation.LOG（自然对数）

Returns:  一个新的layer或None

通过上面6个基础的操作，可以进行大部分的运算。下面通过一个实例展示其用法：

题目：有a,b两个数组，形状为（1，1，4），将a的最后一个维度的前两个数与后两个数做不同的操作后，进行合并，最后在与b相加输出。

import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
TRT_LOGGER = trt.Logger()

def main():
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network:
        builder.max_workspace_size = 1<<30
        builder.max_batch_size = 1
        # 输入层
        input_layer = network.add_input(name="input_layer", dtype=trt.float32, shape=(1, 1, 4))

        # 每个轴都要指定stride
        a1 = network.add_slice(input_layer, (0, 0, 0), (1, 1, 2), (1, 1, 1))
        a2 = network.add_slice(input_layer, (0, 0, 2), (1, 1, 2), (1, 1, 1))

        # a1 的操作
        a1 = network.add_scale(a1.get_output(0), trt.ScaleMode.UNIFORM, \
                          scale=trt.Weights(np.array([2], dtype=np.float32)))
        a1 = network.add_unary(a1.get_output(0), trt.UnaryOperation.EXP)

        # a2 的操作
        a2 = network.add_scale(a2.get_output(0), trt.ScaleMode.UNIFORM, \
                               shift=trt.Weights(np.array([2], dtype=np.float32)))
        # 设置concat, 合并a1, a2
        a = network.add_concatenation([a1.get_output(0),a2.get_output(0)])
        a.axis = 2

        # 数组b
        b = np.ones((1, 1, 4), dtype=np.float32)*2
        b = network.add_constant((1, 1, 4), weights=trt.Weights(b))


        output = network.add_elementwise(a.get_output(0), b.get_output(0), trt.ElementWiseOperation.SUM)
        network.mark_output(output.get_output(0))

        engine = builder.build_cuda_engine(network)
        # Allocate host memory for inputs and outputs.
        h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=np.float32)
        h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=np.float32)

        # Allocate device memory for inputs and outputs.
        d_input = cuda.mem_alloc(h_input.nbytes)
        d_output = cuda.mem_alloc(h_output.nbytes)

        # Create a stream in which to copy inputs/outputs and run inference.
        stream = cuda.Stream()
    with engine.create_execution_context() as context:
          # Transfer input data to the GPU.
          a = np.ones((1, 1, 4))
          print('input:a', a)
          np.copyto(h_input, a.ravel())
          cuda.memcpy_htod_async(d_input, h_input, stream)
          # Run inference.
          context.execute_async(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
          # Transfer predictions back from the GPU.
          cuda.memcpy_dtoh_async(h_output, d_output, stream)
          # Synchronize the stream
          stream.synchronize()
          # Return the host output.
    return h_output

if __name__ == '__main__':
    tr = main()
    tr = np.reshape(tr, (1, 1, 4))
    print('tensorrt:',tr)
    # python验证
    a = np.ones((1, 1, 4))
    b = np.ones((1, 1, 4))*2
    a[:,:,:2] = np.exp(a[:,:,:2]*2)
    a[:, :, 2:] = a[:, :, 2:] + 2
    c = a + b
    print('python:',c)

输出：
input:a [[[1. 1. 1. 1.]]]
tensorrt: [[[9.389056 9.389056 5.       5.      ]]]
python : [[[9.3890561 9.3890561 5.        5.       ]]]

2.7 add_plugin_v2

add_plugin_v2(self: tensorrt.tensorrt.INetworkDefinition, 
            inputs: List[tensorrt.tensorrt.ITensor], 
            plugin: tensorrt.tensorrt.IPluginV2) → tensorrt.tensorrt.IPluginV2Layer
功能：注册插件
Parameters :  input1 - 输入tensor列表，
              plugin - 插件函数

Returns:  一个新的layer或None

除了已经编好的层之外，还有一些特别的插件可以自定义一些操作，官方有写好的插件，也可以自己定义自己的插件。目前主要介绍一些官方的插件：

import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
TRT_LOGGER = trt.Logger()
#加载插件库
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
#获得所有支持的插件
PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list
for plugin_creator in PLUGIN_CREATORS:
    print(plugin_creator.name)

输出:
FancyActivation
ResizeNearest
Split
InstanceNormalization
GridAnchor_TRT
NMS_TRT
Reorg_TRT
Region_TRT
Clip_TRT
LReLU_TRT
PriorBox_TRT
Normalize_TRT
RPROI_TRT
BatchedNMS_TRT
以Clip_TRT和BatchedNMS_TRT作为例子介绍：

Clip_TRT

#使用插件，必须加载插件库
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
###"Clip_TRT"
def get_trt_plugin(plugin_name):
        plugin = None
        for plugin_creator in PLUGIN_CREATORS:
            if plugin_creator.name == plugin_name:
                # 收集参数，各个参数的意义参考https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/c_api/_nv_infer_plugin_8h.html#af308dcae61dab659073bc91c6ba63a7e
                Clip_slope_field = trt.PluginField("clipMin", np.array([1.0], dtype=np.float32), \
                                                    trt.PluginFieldType.FLOAT32)
                Clip_slope_field2 = trt.PluginField("clipMax", np.array([3.0], dtype=np.float32),\
                                                    trt.PluginFieldType.FLOAT32)
                field_collection = trt.PluginFieldCollection([Clip_slope_field,Clip_slope_field2])
                plugin = plugin_creator.create_plugin(name=plugin_name, field_collection=field_collection)
        return plugin


def main():
  with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network:
    builder.max_workspace_size = 1 << 30
    input_layer = network.add_input(name="input_layer", dtype=trt.float32, shape=(5, ))
    # 加入plugins层
    output = network.add_plugin_v2(inputs=[input_layer], plugin=get_trt_plugin("Clip_TRT"))

BatchedNMS_TRT

• 参考https://github.com/NVIDIA/TensorRT/tree/master/plugin/batchedNMSPlugin
• 这个插件可以在GPU上使用加速NMS过程，有两个输入boxes input 和scores input:

boxes input
boxes input的形状shape是[batch_size, number_boxes, number_classes, number_box_parameters],number_box_parameters表示box的定位信息[x_min, y_min, x_max, y_max]。例如，如果你的模型输出在一幅图上产生了8732个候选框，有100个预测类，则boxes input的形状shape为[8732, 100, 4]。
scores input
scores input的形状shape为[batch_size, number_boxes, number_classes],表示各个类别的概率。

• 这个插件会产生4个输出：

1、num_detections : shape=[batch_size, 1], 最后一个维度是一个INT32的标量，表示有效的探测个数，可以少于预设的keepTopK，也表示下面三个对应的输出都是有效的。
2、nmsed_boxes : [batch_size, keepTopK, 4]，NMS后的box坐标。
3、nmsed_scores : [batch_size, keepTopK]，对应的概率值。
4、nmsed_classes : [batch_size, keepTopK]，对应的类别。

• 参数：
• 使用实例：

trt.init_libnvinfer_plugins(TRT_LOGGER, '')
# "BatchedNMS_TRT"
def get_trt_plugin(plugin_name):
        plugin = None
        for plugin_creator in PLUGIN_CREATORS:
            if plugin_creator.name == plugin_name:

                shareLocation = trt.PluginField("shareLocation", np.array([1], dtype=np.int32), \
                                       trt.PluginFieldType.INT32)
                backgroundLabelId = trt.PluginField("backgroundLabelId", np.array([-1], dtype=np.int32),\
                                                    trt.PluginFieldType.INT32)
                numClasses = trt.PluginField("numClasses", np.array([1], dtype=np.int32), \
                                                    trt.PluginFieldType.INT32)
                topK = trt.PluginField("topK", np.array([500], dtype=np.int32), \
                                             trt.PluginFieldType.INT32)
                keepTopK = trt.PluginField("keepTopK", np.array([10], dtype=np.int32), \
                                       trt.PluginFieldType.INT32)
                scoreThreshold = trt.PluginField("scoreThreshold", np.array([0.7], dtype=np.float32),\
                                                    trt.PluginFieldType.FLOAT32)
                iouThreshold = trt.PluginField("iouThreshold", np.array([0.3], dtype=np.float32), \
                                                 trt.PluginFieldType.FLOAT32)
                isNormalized = trt.PluginField("isNormalized", np.array([1], dtype=np.int32), \
                                                trt.PluginFieldType.INT32)
                field_collection = trt.PluginFieldCollection([shareLocation, backgroundLabelId,
                                                              numClasses, topK, keepTopK,
                                                              scoreThreshold, iouThreshold,
                                                              isNormalized])
                plugin = plugin_creator.create_plugin(name=plugin_name, field_collection=field_collection)
        return plugin
output = network.add_plugin_v2(inputs=[location.get_output(0),scores.get_output(0)], \
                            plugin=get_trt_plugin("BatchedNMS_TRT"))
print(output.get_output(1).shape)
network.mark_output(output.get_output(1))

3、将人脸检测的后处理操作加入tensorrt模型中

  通过上面介绍的各种插件，可以把上一节中人脸检测模型的后处理部分整合到一起，一方面可以使模型更简洁方便部署，另一方面也可与后面的人脸识别模型无缝的联结在一起。