6. 自定义算子匹配教程

[Corerain]

自定义算子匹配教程

• 自定义算子匹配教程
  • 1. 为什么需要自定义算子匹配
  • 2.自定义算子匹配的方式
  • 2. Scope Match
    • 2.1 Scope Match 规则介绍
    • 2.2 Scope Match 示例
      • 2.2.1 冻结模型
      • 2.2.2 生成 SG IR
  • 3. Path Match
    • 3.1 Path Match 规则介绍
    • 3.2 Path Match 示例
      • 3.2.1 冻结模型
      • 3.2.2 生成 SG IR
  • 4. 如何量化带有自定义子图的模型

通过本教程，用户能够了解或掌握：

1. 为什么需要自定义算子匹配
2. 自定义算子匹配的方法
3. 带有自定义子图的模型的量化方法

Note：
本教程以虚拟机为基本环境，Linxu用户则需要进入docker。

1. 为什么需要自定义算子匹配

在《CAISA架构与Rainbuilder编译工具介绍》中，我们列出了compiler能够解析的算子列表。虽然compiler能够支持常见的算子，但仍然无法完全满足用户对于其他算子的需求。

为了解决模型中含有compiler无法解析的子图的问题，因此我们在compiler 中提供了自定义算子匹配的功能。

2.自定义算子匹配的方式

我们提供了两种自定义算子匹配方式：

• 域匹配 Scope Match: 利用 TensorFlow 的 scope 特性，匹配同一个 scope 下的节点作为一个子图
• 路径匹配 Path Match: 无法使用 scope 特性时，利用给定起始节点和终止节点来匹配子图

Note:
当匹配到一个自定义子图时，用户需要给这个子图设置一个唯一标识符，不可与 compiler 中的 op 名称重复。

以下是 compiler 中内部使用的 op 名称列表:

输入op	核心op	激活函数op	辅助	常量	预留
Input	ArgMax	LeakyRelu	DataCopy	Const	DropOut
	AvgPool	Relu	DataCrop		Split
	BatchNorm	Relu6	FakeQuant		Unstack
	ConcatV2		FakeMinMax
	Conv2D
	Conv2DBackpropInput
	DepthwiseConv2D
	Flatten
	FullyConnected
	LRN
	MaxPool2D
	Mean
	Pad
	Reshape
	Shortcut
	Softmax
	Upsample

2. Scope Match

2.1 Scope Match 规则介绍

用户需要按照给定的规则来定义配置文件，配置文件使用的是 Json 格式，以下是规则介绍

Parameter	Description
name	规则名称
*scope_regexp	用于描述 scope 的正则表达式, 如 */ROIPooling/*
*target_op_name	自定义子图的 op 名称, 不能与内部的 op 重复
node_attr_map	节点属性到子图属性之间的映射，即提取单个 Const 节点的值

Note: 带 * 的为必填项

2.2 Scope Match 示例

我们先定义一个简单的模型net.py，它只有两层，第一层是compiler 可以解析的卷积，第二层则为暂时无法解析的 l2_normalize。

模型定义如下:

import numpy as np
import tensorflow as tf

slim = tf.contrib.slim

def net_forward(inputs):
  with tf.variable_scope('Corerain'):
    net = slim.conv2d(inputs, 32, [3, 3], scope='conv1')
    net = tf.nn.l2_normalize(net, name="normalize")

  return net


if __name__=='__main__':
  x = tf.placeholder(tf.float32, shape=(1, 256, 256, 3), name='x')
  y = net_forward(x)

  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver = tf.train.Saver()
    saver.save(sess, './tmp/test.ckpt')

执行 python net.py，保存模型的 checkpoint，因为我们仅仅是作为演示，不需要训练此模型。

2.2.1 冻结模型

获得ckpt文件后，需要通过compiler冻结模型。此步骤与“Compiler与网络部署指南”中“冻结模型”一步中一致。

$ RbCli freeze ./tmp -m ./tmp/model.pb --logdir ./tmp
=================================================================
*    ______ _     _____                       _ _               *
*    | ___ \ |   /  __ \                     (_) |              *
*    | |_/ / |__ | /  \/ ___  _ __ ___  _ __  _| | ___ _ __     *
*    |    /| '_ \| |    / _ \| '_ ` _ \| '_ \| | |/ _ \ '__|    *
*    | |\ \| |_) | \__/\ (_) | | | | | | |_) | | |  __/ |       *
*    \_| \_|_.__/ \____/\___/|_| |_| |_| .__/|_|_|\___|_|       *
*                                      | |                      *
*                                      |_|                      *
=================================================================
The version of RbCompiler installed is 1.4.1, TensorFlow installed is: 1.12.0
Please report any bug or issue to <vincent.zhao@corerain.com>

You haven't specified any output node, please select one or many from the following list:
[  0] Corerain/normalize
Please enter the indices of output nodes, separated by ','. If you want to select all, please enter 'all': all
Select output node names: ['Corerain/normalize']
Freezing ./ to model.pb ...

在冻结模型之前，可以先生成 tensorboard 的日志文件，方便我们可视化模型：

$ tensorboard --logdir ./tmp --host 127.0.0.1 --port 6006

在浏览器上打开 http://127.0.0.1:6006, 你可以看到如何模型的结构。

双击无法识别的 normalize:

在这之前，我们需要先观察不能识别的子图的结构，以确定使用何种方式进行匹配，经过观察，normalize 下面的节点都有相同的 Corerain/normalize scope，我们可以采用 Scope Match 匹配规则

Json 配置文件如下(layer.json)

[{
  "name": "l2_normalize",
  "scope_regexp": "Corerain/normalize",
  "target_op_name": "l2_normalize",
  "node_attr_map": [{
    "const_node_name": "Corerain/normalize/Maximum/y",
    "data_field_name": "eplison"
  }]
}]

Note:
const_node_name 是节点的名称, data_field_name 是序列化的键名称

2.2.2 生成 SG IR

RbCli sg ./tmp/model.pb -c layer.json

生成的 SG IR 如下所示(... 表示省略了 Const 和 Conv2D 节点)

name: "model"
version: "v2"
node {
  name: "x"
  op: "Input"
  device: CPU
  type: T_FLOAT
  attr_map {
    key: "data_format"
    value {
      s: "0,2,3,1"
    }
  }
  attr_map {
    key: "shape"
    value {
      list {
        i: 1
        i: 256
        i: 256
        i: 3
      }
    }
  }
}
....
....
node {
  name: "Corerain/conv1/Relu"
  input: "Corerain/conv1/Conv2D"
  op: "Relu"
  device: CPU
  type: T_FLOAT
  attr_map {
    key: "type"
    value {
      s: "Relu"
    }
  }
}
node {
  name: "Corerain/normalize/Square"
  input: "Corerain/conv1/Relu"
  op: "l2_normalize"
  device: CPU
  type: T_FLOAT
  attr_map {
    key: "eplison"
    value {
      f: 9.999999960041972e-13
    }
  }
  attr_map {
    key: "is_output"
    value {
      b: true
    }
  }
}

3. Path Match

3.1 Path Match 规则介绍

Parameter	Description
name	规则名称
*start	起始节点的名称列表
*end	终止节点的名称列表
*target_op_name	提取的自定义 SG 节点的操作名称, 不能与内部的 op 重复
node_attr_map	节点属性到子图属性之间的映射，即提取单个节点的值

Note: 带 * 的是必填项

Path Match规则只是把 scope_regexp 替换成了 start 和 end 参数，start 和 end 是一个列表，它表示子图的起始节点，但是起始节点本身不包括在子图中，end 节点是终止节点，终止节点包括在子图里，也就是说这是一个 左开右闭 (start, end] 的区间。

为什么是左开右闭区间？

因为在提取子图的节点时，我们不光要提取节点本身，还要提取节点的 Const 输入，但是有些输入它不是直接连接的 Const，它可能会经过 Const -> pack -> ... 等操作，变成一个非 Const 直接作为节点的输入，导致无法判断这个输入是否是上一层的节点，所以起始节点不应该包括在待提取的子图中。

如：A -> B -> C -> D -> F， 我们定义区间 (A, D] 则表示提取 B C D 节点成为一个子图。

3.2 Path Match 示例

仍然以net.py模型为例，只是使用 Path Match 规则，首先我们要确定 normalize 的起始节点。

从图中可知，起始节点是 Relu， 全名为 Corerain/conv1/Relu, 终止节点是 Corerain/normalize/(normalize)，实际在文件中读取的名称是 Corerain/normalize。

Json 配置文件如下(layer.json)

[{
  "name": "l2_normalize",
  "start": ["Corerain/conv1/Relu"],
  "end": ["Corerain/normalize"],
  "target_op_name": "l2_normalize",
  "node_attr_map": [{
    "const_node_name": "Corerain/normalize/Maximum/y",
    "data_field_name": "eplison"
  }]
}]

3.2.1 冻结模型

同样地，获得ckpt文件后，需要通过compiler冻结模型。

RbCli freeze ./tmp -m ./tmp/model.pb --logdir ./tmp

3.2.2 生成 SG IR

$ RbCli sg ./tmp/model.pb -c layer.json

生成的 SG IR 如下所示(... 表示省略了 Const 和 Conv2D 节点)

name: "model"
version: "v2"
node {
  name: "x"
  op: "Input"
  device: CPU
  type: T_FLOAT
  attr_map {
    key: "data_format"
    value {
      s: "0,2,3,1"
    }
  }
  attr_map {
    key: "shape"
    value {
      list {
        i: 1
        i: 256
        i: 256
        i: 3
      }
    }
  }
}
....
....
node {
  name: "Corerain/conv1/Relu"
  input: "Corerain/conv1/Conv2D"
  op: "Relu"
  device: CPU
  type: T_FLOAT
  attr_map {
    key: "type"
    value {
      s: "Relu"
    }
  }
}
node {
  name: "Corerain/normalize"
  input: "Corerain/conv1/Relu"
  op: "l2_normalize"
  device: CPU
  type: T_FLOAT
  attr_map {
    key: "eplison"
    value {
      f: 9.999999960041972e-13
    }
  }
  attr_map {
    key: "is_output"
    value {
      b: true
    }
  }
}

4. 如何量化带有自定义子图的模型

在量化之前，用户需要实现自定义子图的计算代码，此部分由 C++ 实现，并且根据上面的步骤实现自定义子图算子。

仍以net.py为例，首先需要查阅 tf.nn.l2_normalize() 的计算原理。 TensorFlow 官方给出的公式如下：

output = x / sqrt(max(sum(x**2), epsilon))

创建C++ 文件l2_normalize.cc以实现相关功能:

/*
Normalizes along dimension axis using an L2 norm
Author: zouwei
Date: 2019/3/27
*/

#include "RbRuntime/core.h"

#include <iostream>
#include <memory>
#include <math.h>
#include <algorithm>

using namespace RbRuntime::tensor;
using namespace RbRuntime::op_impl;

namespace RbRuntime {
namespace nn {
namespace ops {

template <typename Device, typename T> class L2NormalizationOpImpl : public OpImpl {
public:
  explicit L2NormalizationOpImpl(OpImplContext *ctx) : OpImpl(ctx) {}

  void Run(OpImplEnv *env) {
    auto node_def = this->ctx_->node()->node_def();
    auto epsilon = node_def.attr_map().find("epsilon")->second.f();

    auto input = env->input(0);
    auto output = env->output(0);

    auto in = input->tensor<T, 4>();
    auto out = output->tensor<T, 4>();

    auto total = 0.0f;
    for (int i = 0; i < input->num_elems(); i++) {
      total += static_cast<float>(in[i] * in[i]);
    }

    auto val = sqrt(std::max(total, epsilon));
    for (int i = 0; i < input->num_elems(); i++) {
      out[i] = static_cast<float>(in[i]) / val;
    } 

  }

  TensorShape GetOutputTensorShape(OpImplEnv *env) {
    return env->input(0)->shape();
  }
};

INIT_BY_ALL_TYPE_W(REGISTER_OP_IMPL_ON_CPU, "L2Normalization", L2NormalizationOpImpl)
} // namespace ops
} // namespace nn
} // namespace RbRuntime

对应的CMakeList.txt如下:

cmake_minimum_required(VERSION 3.5)
project(l2_normalize)

set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -ldl")

add_library(l2_normalize SHARED l2_normalize.cc)
target_include_directories(l2_normalize PUBLIC /usr/local/include)

target_link_libraries(l2_normalize /usr/local/lib/libRbRuntime.so)

/usr/local/lib/libRbRuntime.so 为用户安装的 Runtime 库。（虚拟机、linux docker均为此路径）

执行：

mkdir build && cd build
cmake ..
make -j2

把生成的 libl2_normalize.so 放到和 *_sg.pbtxt 同一目录下，即可执行有自定义子图的 RbCli quant 命令。