跳转至

NPU编译器使用*

在使用 NPU 编译器gxnpuc前,请仔细阅读下面两篇技术文档:

gxnpuc用于将开源框架的网络模型文件转换为适配国芯 NPU 处理器的离线模型文件;

1. gxnpuc 工具链功能及对应参数简述*

1.1 常规功能参数*

--help

  • 参数说明:

    打印 gxnpuc 工具链的相关参数说明

  • 使用示例

    $ gxnpuc --help
    
    usage: gxnpuc [-h] [--cmpt] [--list] [-c {LEO,APUS,GRUS,V100,V120,V150}]
                  [-f {TF,PT}] [-V] [-v] [-m] [-w] [-s] [-q]
                  [config_filename]
    
    NPU Compiler
    
    positional arguments:
      config_filename       config file
    
    optional arguments:
      -h, --help            show this help message and exit
      --cmpt                get version compatibility information between npu-
                            core, python, and frameworks
      --list                list supported ops
      -c {LEO,APUS,GRUS,V100,V120,V150}, --core_name {LEO,APUS,GRUS,V100,V120,V150}
                            subparameter of --list, specify NPU Core for listing
                            supported ops
      -f {TF,PT}, --framework {TF,PT}
                            subparameter of --list, specify Deep Learning
                            Framework for listing supported ops
      -V, --version         show program's version number and exit
      -v, --verbose         verbosely list the processed ops
      -m, --meminfo         verbosely list memory info of ops
      -w, --weights         print compressed weights (GRUS only)
      -s, --save_hist       save histograms of weights value to 'npu_jpgs'
                            directory (GRUS only)
      -q, --quant           inference and generate quant file
    

--version

  • 参数说明:

    显示当前编译器版本信息

  • 使用示例

    gxnpuc --version
    

--list

  • 参数说明:

    列出当前NPU编译器支持的算子信息

  • 关联子参数

    子参数仅在 --list 参数使用时有效,且无强制要求使用

    -c

    指定芯片版本,参数值可选范围: {LEO,APUS,GRUS,V100,V120,V150}
    

    -f

    指定前端深度学习框架,参数值可选范围: {TF、PT}
    
  • 使用示例:

    列出当前 NPU 编译器支持的全部算子信息

    gxnpuc --list
    

    列出当前 NPU GRUS 编译器支持的全部算子信息

    gxnpuc --list -c GRUS
    

    列出当前 NPU GRUS 编译器支持的全部 PyTorch 算子信息

    gxnpuc --list -c GRUS -f PT
    

--cmpt

  • 参数说明:

    打印编译器所支持的Python、芯片型号、前端DL框架版本信息,以及三者之间的兼容关系

  • 使用示例:

    gxnpuc --cmpt
    

1.2 模型编译转换参数*

config_filename

  • 参数说明:

    指定编译配置文件的路径以及文件名,模型转换及模型调优时会读取该配置文件

  • 关联子参数(仅在 config_filename 正确配置时有效)

    -v

    功能描述: 编译模型时,打印 NPU 模型结构信息
    

    -m

    功能描述: 编译模型时,打印 NPU 模型中各算子节点使用的内存状态
    

    -w

    功能描述: 打印 NPU 模型中压缩的权重信息(仅GRUS支持)
    

    -s

    功能描述: 保存模型权重直方分布图至npu_jpgs文件夹下(仅GRUS支持)
    
  • 使用示例:

    NPU GRUS 编译器以 config.yaml 为原始模型的转换配置文件,同时开启所有子参数功能

    gxnpuc config.yaml -v -m -w -s
    

gxnpuc功能参数使用注意事项

上述 --version/-V, --list, --cmpt, config_filename 四组参数互斥,不能同时使用

2. 模型编译转换配置文件说明*

2.1 TensorFlow 配置项*

配置项 参数取值 参数说明
CORENAME GRUS 芯片型号
NPU_UNIT NPU32 指定 NPU 型号
FRAMEWORK TF 指定待转换模型的前端DL框架类型
MODEL_FILE 模型文件名及路径 例: ./model.pb 指定待转换模型的文件名及路径
OUTPUT_TYPE c_code 指定编译器输出 NPU 文件的格式
OUTPUT_FILE NPU文件名 例: npu.h 指定编译器输出 NPU 文件的名称
INPUT_OPS op_name: shape 指定 NPU 模型所有输入节点的相关信息
OUTPUT_OPS [output_name, ...] 指定 NPU 模型所有输出节点的相关信息
FP16_OUT_OPS [out_state_name, ...] 指定 NPU 模型以 Float16 格式输出的节点
FUSE_BN true / false (默认值 false) 是否使能 BN 参数融合功能
COMPRESS true / false (默认值 false) 是否使能全连接权重量化压缩功能
CONV2D_COMPRESS true / false (默认值 false) 是否使能卷积权重量化压缩功能
EXCLUDE_COMPRESS_OPS [weight_op_name, ...] 指定部分权重节点可不进行量化压缩
WEIGHT_MIN_MAX weight_op_name: [min, max] 指定权重节点量化压缩时的最小值和最大值
WEIGHT_CACHE_SIZE 具体分配的内存值 例: 10240 指定分配 SRAM 中存储权重的内存大小

注意事项

  • 待转换的原始模型必须为 FrozenPB 格式

INPUT_OPS

  • 参数格式

    op_name: shape

    op_name —— 模型输入节点的名称
    
    shape   —— 对应节点名为 op_name 的输入在推理时的 shape
    
  • 使用示例

    TensorFlow 框架中搭建模型时,通常定义占位符作为计算图的输入,用户需要为占位符指定具体的标识符作为输入名称

    本例代码中,模型定义了四组占位符(模型输入),并分别配置 Feats、State_c0、State_c1、State_c2 标识符作为输入名称

    其中 state0_in、state1_in、state2_in 为上一帧输出状态值(初始帧时为全零张量)

    inputs    = tf.placeholder(tf.float32, [1, 1, 64], name="Feats")
    state0_in = tf.placeholder(tf.float32, [1, 3, 64], name="State_c0")
    state1_in = tf.placeholder(tf.float32, [1, 4, 64], name="State_c1")
    state2_in = tf.placeholder(tf.float32, [1, 5, 64], name="State_c2")
    

    故该模型配置文件中的 INPUT_OPS 参数可按如下方式配置:

    config.yaml
    INPUT_OPS:
        Feats:    [1, 1, 64]
        State_c0: [1, 3, 64]
        State_c1: [1, 4, 64]
        State_c2: [1, 5, 64]
    

OUTPUT_OPS

  • 参数格式

    [output_name, ...]

    output_name 为模型输出节点的名称
    
  • 使用示例

    TensorFlow 框架中搭建模型时,为方便 NPU 的编译配置,用户可以通过 tf.identity 接口对输出张量实现拷贝并重命名的操作

    本例代码中,为四组输出张量分别配置 Result、State_c0_out、State_c1_out、State_c2_out 标识符作为输出名称

    其中 state0_out、state1_out、state2_out 为当前一帧的输出状态值

    outputs, states = fsmn_layer(...)
    
    result_out = tf.identity(outputs,   name="Result")
    state1_out = tf.identity(states[0], name="State_c0_out")
    state2_out = tf.identity(states[1], name="State_c1_out")
    state3_out = tf.identity(states[2], name="State_c2_out")
    

    故该模型配置文件中的 INPUT_OPS 参数可按如下方式配置:

    config.yaml
    OUTPUT_OPS: [State_c0_out, State_c1_out, State_c2_out, Result]
    

    注意事项

    • 输出状态结点必须放在预测输出节点之前

FP16_OUT_OPS

  • 参数相关概述

    NPU 内部运算采用 FP16 格式数据,输入输出张量都是 FP16 格式;

    NPU 支持格式转换 FP16_TO_FP32 功能,不支持格式转换 FP32_TO_FP16 功能;

    循环神经网络的处理流程中,在每个时间步下,网络会接收当前时刻的输入以及前一时间步的隐藏状态,通过这些信息生成当前时刻的隐藏状态和相应的预测结果;

    在实际应用中,预测结果对应的输出张量会先转换为 FP32 格式后,再进行预测处理,隐藏状态对应的输出张量则直接以 FP16 格式作为下一帧的输入;

    用户需要按照具体模型结构,将输出节点的数据格式划分为 FP32 和 FP16;

    用户可通过 FP16_OUT_OPS 参数指定哪些输出节点作为 FP16 格式输出;

  • 参数格式

    [output_state_name, ...]

    output_state_name 为模型隐藏状态对应的输出节点名称
    
  • 使用示例

    沿用 OUTPUT_OPS 使用示例

    该模型配置文件中的 FP16_OUT_OPS 参数可按如下方式配置:

    config.yaml
    FP16_OUT_OPS: [State_c0_out, State_c1_out, State_c2_out]
    

EXCLUDE_COMPRESS_OPS

  • 参数相关概述

    NPU 编译器权重量化功能开启后,若模型推理性能较差,用户可根据权重分布直方图,分析各权重节点是否量化友好,并可指定哪些权重节点不进行量化处理

  • 参数格式

    [weight_op_name, ...]

    weight_op_name 为需要关闭量化处理的权重节点名称
    

  • 使用示例

    本例中,假设模型中部分卷积权重量化处理后,导致整体模型性能较差,故不对其进行量化处理

    配置文件中的 EXCLUDE_COMPRESS_OPS 参数可按如下方式配置:

    config.yaml
    EXCLUDE_COMPRESS_OPS: [conv2d_5/Conv2D/ReadVariableOp/_74__cf__74,
                           conv2d_6/Conv2D/ReadVariableOp/_75__cf__75]
    

WEIGHT_MIN_MAX

  • 参数相关概述

    NPU 编译器采用训练后量化(PTQ, Post Train Quantization),量化策略选择 MinMax 方法统计权重数据分布范围;

    NPU 编译器权重量化功能开启后,若模型推理性能较差,用户可根据权重分布直方图,分析各权重节点数据分布范围的统计是否合理,并可通过该参数直接配置指定权重节点的数据分布范围;

  • 参数格式

    weight_op_name: [min, max]

    weight_op_name —— 模型指定某个权重节点的名称
    
    min, max       —— 指定 weight_op_name 张量的最大最小值
    

  • 使用示例

    本例中指定两组卷积的量化处理范围,配置文件中的 EXCLUDE_COMPRESS_OPS 参数可按如下方式配置:

    config.yaml
    WEIGHT_MIN_MAX:
        conv2d_5/Conv2D/ReadVariableOp/_74__cf__74: [min0, max0]
        conv2d_6/Conv2D/ReadVariableOp/_75__cf__75: [min1, max1]
    

EXCLUDE_COMPRESS_OPS、WEIGHT_MIN_MAX 参数的具体使用场景和使用方式请参考 NPU 量化精度调试

2.2 PyTorch 配置项*

注意事项

若用户需要编译转换 PyTorch 模型,则 NPU 编译器版本至少 1.6.0b0 以上,且必须使用 Python3.7

配置项 参数取值 参数说明
CORENAME GRUS 芯片型号
NPU_UNIT NPU32 指定 NPU 型号
FRAMEWORK PT 指定待转换模型的前端DL框架类型
MODEL_FILE 模型文件名及路径 例: ./model.pb 指定待转换模型的文件名及路径
OUTPUT_TYPE c_code 指定编译器输出 NPU 文件的格式
OUTPUT_FILE NPU文件名(例 npu.h) 指定编译器输出 NPU 文件的名称
INPUT_OPS input_index: shape 指定 NPU 模型所有输入节点的相关信息
INPUT_NCX_TO_NXC [input_index, ...] 是否转换 NPU 模型输入张量的数据排布格式
FP16_OUT_OPS [out_state_index, ...] 指定 NPU 模型中以 Float16 格式输出的节点
FUSE_BN true / false (默认值 false) 是否使能 BN 参数融合功能
COMPRESS true / false (默认值 false) 是否使能全连接权重量化压缩功能
CONV2D_COMPRESS true / false (默认值 false) 是否使能卷积权重量化压缩功能
EXCLUDE_COMPRESS_OPS [weight_op_name, ...] 指定部分权重节点可不进行量化压缩
WEIGHT_MIN_MAX weight_op_name: [min, max] 指定权重节点量化压缩时的最小值和最大值
WEIGHT_CACHE_SIZE 具体分配的内存值
(例: 10240)
指定分配 SRAM 中存储权重的内存大小

注意事项

  • 待转换的原始模型必须为 jit.ScriptModule 格式

  • 待转换模型文件 必须与配置文件在同一路径下

  • 编译器1.6.0.b0版本以上,Python3.7,PyTorch 1.10 - 1.13

INPUT_OPS

  • 参数格式

    input_index: shape

    input_index —— 输入节点的索引
    
    shape       —— 对应输入在 NPU 模型中的推理 shape
    
  • 使用示例

    由于 PyTorch 框架中采用动态计算图,各算子节点 name 是自动生成的,导致无法为输入配置标识符;

    因此 NPU 编译器转换 PyTorch 模型时,配置文件中 INPUT_OPS 参数采用索引值进行输入张量的映射;

    本例中,以继承 nn.Module 基类的方式来构建自定义 PyTorch 模型,并简述 INPUT_OPS 参数的配置方式:

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3)
            self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.conv2 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
            self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.adaptive_pool = nn.AdaptiveMaxPool2d((1,1))
            self.flatten = nn.Flatten()
            self.linear1 = nn.Linear(64,32)
            self.relu    = nn.ReLU()
            self.linear2 = nn.Linear(32,1)
    
        def forward(self, x, y):
            x = self.conv1(x)
            x = self.pool1(x)
            y = self.conv2(y)
            y = self.pool2(y)
            z = torch.concat([x,y], dim=1)
            z = self.adaptive_pool(z)
            z = self.flatten(z)
            z = self.linear1(z)
            z = self.relu(z)
            y = self.linear2(z)
            return y
    
    net           = Net()
    input0_tensor = torch.randn([1, 3, 32, 32])
    input1_tensor = torch.randn([1, 1, 32, 32])
    
    output_tensor = net(input0_tensor, input1_tensor)
    

    由 forward 方法可知上述模型定义了两组输入 x、y,输入 x 的索引为0,输入 y 的索引为1;

    故该模型配置文件中的 INPUT_OPS 参数可按如下方式配置:

    config.yaml
    INPUT_OPS:
        0: [1, 3, 32, 32]
        1: [1, 1, 32, 32]
    
  • 约束条件

    PyTorch 模型转换时,需要保证输入是张量,不可以是张量列表或张量元组

    NPU 编译器不支持上述类型的输入格式,用户需要将张量列表和张量元组拆分为张量进行输入

    本例模型中,输入采用了张量列表/张量元组(与使用示例中的模型仅在输入有差异)

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3)
            self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.conv2 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
            self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
            self.adaptive_pool = nn.AdaptiveMaxPool2d((1,1))
            self.flatten = nn.Flatten()
            self.linear1 = nn.Linear(64,32)
            self.relu    = nn.ReLU()
            self.linear2 = nn.Linear(32,1)
    
        def forward(self, xy):
            x = self.conv1(xy[0])
            x = self.pool1(x)
            y = self.conv2(xy[1])
            y = self.pool2(y)
            z = torch.concat([x,y], dim=1)
            z = self.adaptive_pool(z)
            z = self.flatten(z)
            z = self.linear1(z)
            z = self.relu(z)
            y = self.linear2(z)
            return y
    
    net           = Net()
    input0_tensor = torch.randn([1, 3, 32, 32])
    input1_tensor = torch.randn([1, 1, 32, 32])
    list_tensor   = [input0_tensor, input1_tensor]
    tuple_tensor  = (input0_tensor, input1_tensor)
    
    # 输入是张量列表
    output_tensor = net(list_tensor)
    
    # 输入是张量元组
    output_tensor = net(tuple_tensor)
    

INPUT_NCX_TO_NXC

  • 参数相关概述

    在深度学习领域中,模型算子节点之间通常采用多维度张量进行数据传输,比如卷积神经网络的特征图通常采用四维度张量存储

    四维度张量的各维度可分别表示为 N: batch; H: height; W: width; C: channels

    由于数据在内存中是以线性排布方式存储,当数据维度的访问顺序变化时,数据在内存中的分布也存在差异; PyTorch 框架中采用 NCHW 顺序,TensorFlow 框架中采用 NHWC 顺序,上述两种访问顺序可称为数据的排布格式(data format)

    NPU 计算密集型算子采用 NHWC、NLC 数据排布格式进行存储,若 NPU 计算密集型算子的输入张量是 NCHW 或 NCL 数据排布格式,则编译器会在该算子前插入转置节点进行数据格式的转换

  • 参数说明

    当原始模型输入张量的数据格式为 NCHW 或 NCL 时,可通过配置该参数,决定 NPU 模型推理时,该输入张量是保持原始 PyTorch 模型中的数据排布格式,还是需要已经 NCHW -> NHWC 或 NCL -> NLC 转换处理后的数据排布格式

  • 参数格式

    [input_index, ...]

    input_index —— 需要进行数据排布格式转换的输入节点的索引
    

    注意事项

    当使能某个输入张量进行数据排布格式转换时,需要同步调整参数 INPUT_OPS 中对应输入的 shape,具体使用可见使用示例

  • 使用示例

    import torch.nn as nn
    
    class Model(torch.nn.Module):
        def __init__(self):
            super(Model, self).__init__()
            self.gru  = torch.nn.GRU(128, 128, batch_first=True, bias=True)
            self.conv = torch.nn.Conv1d(128, 128, 1, 1)
    
        def forward(self, x, h):
            x = self.conv(x)
            x = x.permute(0, 2, 1)
            z = self.gru(x, h)
    
            return z
    
    model        = Model()
    batch        = 1
    seq_length   = 32
    channel      = 128
    input_tensor = torch.randn([batch, channel, seq_length])
    input_state  = torch.randn([1, batch, channel])
    
    output_tensor, output_state = model(input_tensor, input_state)
    

    上述模型在 PyTorch 框架下推理时,input_state 数据排布格式不是 NCL,故无法进行数据排布格式转换;

    input_tensor 数据排布格式为 NCL,故可以选择是否进行数据排布格式转换后,再提供至 NPU 模型输入;

    下面将分别描述格式转换参数使能与不使能的两种情况:

    • 使能参数,输入张量 input_tensor 进行格式转换

      input_tensor 原始模型输入 shape 为 [1, 128, 32],使能格式转换功能后,NPU 模型需要的实际 shape 为 [1, 32, 128]

      模型配置文件的 INPUT_NCX_TO_NXC 和 INPUT_OPS 参数可按如下方式配置:

      config.yaml
      INPUT_NCX_TO_NXC: [0]
      
      INPUT_OPS:
          0: [1, 32, 128]
          1: [1, 1, 128]
      

      下图为模型结构图:

    • 不使能参数,输入张量 input_tensor 保持原始数据排布格式

      模型配置文件的 INPUT_NCX_TO_NXC 和 INPUT_OPS 参数可按如下方式配置:

      config.yaml
      INPUT_NCX_TO_NXC: []
      
      INPUT_OPS:
          0: [1, 128, 32]
          1: [1, 1, 128]
      

      下图为模型结构图:

      不使能 INPUT_NCX_TO_NXC

    对比上述两 NPU 模型结构图可知,格式转换功能开启后,可将输入侧插入的转置节点优化

    该参数能够根据用户需求,更加灵活的设置输入张量的数据格式

FP16_OUT_OPS

  • 参数相关概述

    NPU 内部运算采用 FP16 格式数据,输入输出张量都是 FP16 格式;

    NPU 支持格式转换 FP16_TO_FP32 功能,不支持格式转换 FP32_TO_FP16 功能;

    循环神经网络的处理流程中,在每个时间步下,网络会接收当前时刻的输入以及前一时间步的隐藏状态,通过这些信息生成当前时刻的隐藏状态和相应的预测结果;

    在实际应用中,预测结果对应的输出张量会先转换为 FP32 格式后,再进行预测处理,隐藏状态对应的输出张量则直接以 FP16 格式作为下一帧的输入;

    用户需要按照具体模型结构,将输出节点的数据格式划分为 FP32 和 FP16;

    用户可通过 FP16_OUT_OPS 参数指定哪些输出节点作为 FP16 格式输出;

  • 参数取值

    [out_state_index, ...]

    out_state_index —— 模型隐藏状态对应输出节点的索引
    
  • 使用示例

    PyTorch 多输出模型

    import torch.nn as nn
    
    class TestModel(torch.nn.Module):
        def __init__(self):
            super(TestModel, self).__init__()
            self.gru  = nn.GRU(32, 32, batch_first=True)
            self.conv = nn.Conv1d(32, 32, 1, 1)
    
        def forward(self, x, h):
            t = torch.split(x, [1, 2, 3, 4], dim=1)
    
            y = self.gru(x, h)
            z = torch.sigmoid(y[0])
            z = z.permute(0, 2, 1)
            o = self.conv(z)
    
            return t, y, o
    
    net          = TestModel()
    input_tensor = torch.randn([1, 10, 32])
    state_tensor = torch.randn([1, 1, 32])
    
    output_t, output_y, output_o = model(input_tensor, input_state)
    

    由 forward 方法可知,示例中的模型定义了三组输出 t、y、o,其中输出 t、y 都是张量列表,输出 o 是张量

    NPU 编译器会将张量列表输出展开为张量输出,对于上述模型,NPU 编译器将认为有 7 个输出 t[0], t[1], t[2], t[3], y[0], y[1], o

    模型输出 y[1] 是 gru 模块的隐藏状态对应的输出节点 state,需要用于下一帧推理时输入至 gru 模块,该输出不进行 FP16 -> FP32 转换

    该模型配置文件中的 FP16_OUT_OPS 参数可按如下方式配置:

    config.yaml
        FP16_OUT_OPS: [5]
    

3. 编译模型*

3.1 模型文件准备*

NPU编译器严格限制模型文件格式,不同框架需按照指定方式导出。

TensorFlow

  • 准备 TensorFlow 生成的 CKPT 和 PB 文件,或 saved_model 方式生成的模型文件。
  • 通过 TensorFlow 提供的freeze_graph.py脚本生成 FROZEN_PB 文件。

PyTorch

  • 模型训练完成后,导出模型权重文件
  • 构建 PyTorch 推理模型脚本,并生成 PyTorch Module 实例
  • 转换自定义 PyTorch Module 为 Torch ScriptModule,并将 ScriptModule 实例序列化输出为编译器所需的模型文件

    上述步骤具体可参考 PyTorch 模型转换示例

3.2 编写配置文件*

  • 编写 yaml 配置文件,包括 模型文件名,输出文件名,输出文件类型,是否压缩,输入节点名和维度信息,输出节点名等

3.3 编译生成模型文件*

编译命令如下

$ gxnpuc config.yaml

注意

NPU 工具链编译处理不同深度框架的模型文件时,必须安装对应框架在 NPU 工具链的运行环境 生成的模型文件格式说明请阅读:NPU模型格式说明

4. 一些OP的说明*

4.1 Softmax*

NPU 无法直接支持 Softmax,但在一定条件下,可以通过修改模型让 NPU 支持 Softmax 计算

需要满足条件:

  • Softmax 的输入张量必须为 2 维度,且 batch size 等于 1

模型中的 Softmax 需要改成如下函数:

TensorFlow

def factorize(n):
    for i in range(1, 16):
        if n % i == 0 and n // i <= 15:
            return (n // i, i)
    return ()

def split_and_factorize(n):
    result = []
    while not factorize(n):
        for i in range(n-1, 0, -1):
            if factorize(i):
                result.append(i)
                n -= i
                break
    result.append(n)
    return result

def npu_softmax(x, name=None):
    """ NPU Softmax
    Args:
      x: A non-empty `Tensor`.
      name: A name for the operation (optional).
    Returns
      A `Tensor`.
    """

    # x' = x - max(x)
    # y = exp(x') / sum(exp(x'))
    assert len(x.shape) == 2 and x.shape[0] == 1
    partitions = split_and_factorize(x.shape[1])
    if len(partitions) == 1:
        a, b = factorize(partitions[0])
        pool_shape = [1, a, b, 1]
        x = tf.reshape(x, pool_shape)
        max_ = tf.nn.max_pool(x, ksize=pool_shape, strides=pool_shape, padding='VALID')
    else:
        cnt = 0
        tmp_max_list = []
        for p in partitions:
            a, b = factorize(p)
            pool_shape = [1, a, b, 1]
            tmp_x = tf.reshape(x[:,cnt:cnt+p], pool_shape)
            tmp_max = tf.nn.max_pool(tmp_x, ksize=pool_shape, strides=pool_shape, padding='VALID')
            tmp_max_list.append(tmp_max)
            cnt += p
        tmp_max_len = len(tmp_max_list)
        assert tmp_max_len <= 15
        max_ = tf.concat(tmp_max_list, axis=2)
        max_ = tf.nn.max_pool(max_, ksize=[1,1,tmp_max_len,1], strides=[1,1,tmp_max_len,1], padding='VALID')
    x = tf.reshape(x, [-1, 1])
    x = tf.math.subtract(x, max_)
    exp_x = tf.exp(x)
    exp_x = tf.reshape(exp_x, [1, -1])
    return tf.math.divide(exp_x, tf.math.reduce_sum(exp_x, axis=-1), name=name)

重要

训练用 tensorflow 的 softmax, 训练结束导出 CKPT 前改成 npu_softmax。

PyTorch

def factorize(n):
    for i in range(1, 16):
        if n % i == 0 and n // i <= 15:
            return (n // i, i)
    return ()

def split_and_factorize(n):
    result = []
    while not factorize(n):
        for i in range(n-1, 0, -1):
            if factorize(i):
                result.append(i)
                n -= i
                break
    result.append(n)
    return result

def npu_softmax(x, partitions):
    """ NPU Softmax
    Args:
      x: A non-empty `Tensor`. limitation: len(x.shape) == 2 and x.shape[0] == 1
      partitions: A list of integers specifying the partition sizes for the softmax operation.
    Returns
      A `Tensor` representing the real softmax output.
    """

    if len(partitions) == 1:
        a, b = factorize(partitions[0])
        pool_shape = [1, 1, a, b]

        x_ = x.reshape(pool_shape)
        max_ = torch.nn.functional.max_pool2d(x_, (a,b), (a,b))
    else:
        cnt = 0
        tmp_max_list = []
        for p in partitions:
            a, b = factorize(p)
            pool_shape = [1, 1, a, b]
            x_  = x[:,cnt:cnt+p]

            tmp_x = x_.reshape(pool_shape)

            tmp_max = torch.nn.functional.max_pool2d(tmp_x, (a,b), (a,b))
            tmp_max_list.append(tmp_max)

            cnt += p

        tmp_max_len = len(tmp_max_list)

        max_ = torch.concat(tmp_max_list, axis=3)
        max_ = torch.nn.functional.max_pool2d(max_, (1,tmp_max_len), (1,tmp_max_len))

    max_  = max_.squeeze(0).squeeze(0)
    x     = x - max_
    exp_x = torch.exp(x)
    sum_x = torch.sum(exp_x, dim=-1)
    return exp_x / sum_x

重要

训练用 pytorch 的 softmax, 构建推理脚本时改成 npu_softmax。

4.2 LogSoftmax*

NPU 无法直接支持 LogSoftmax,但在一定条件下,可以通过修改模型让 NPU 支持 LogSoftmax 计算

需要满足条件:

  • LogSoftmax 的输入张量必须为 2 维度,且 batch size 等于 1

模型中的 LogSoftmax 需要改成如下函数:

TensorFlow

def factorize(n):
    for i in range(1, 16):
        if n % i == 0 and n // i <= 15:
            return (n // i, i)
    return ()

def split_and_factorize(n):
    result = []
    while not factorize(n):
        for i in range(n-1, 0, -1):
            if factorize(i):
                result.append(i)
                n -= i
                break
    result.append(n)
    return result

def npu_log_softmax(x, name=None):
    """ NPU LogSoftmax
    Args:
      x: A non-empty `Tensor`.
      name: A name for the operation (optional).
    Returns
      A `Tensor`.
    """
    # x' = x - max(x)
    # y  = x' - log(sum(exp(x')))
    assert len(x.shape) == 2 and x.shape[0] == 1
    partitions = split_and_factorize(x.shape[1])
    if len(partitions) == 1:
        a, b = factorize(partitions[0])
        pool_shape = [1, a, b, 1]
        x = tf.reshape(x, pool_shape)
        max_ = tf.nn.max_pool(x, ksize=pool_shape, strides=pool_shape, padding='VALID')
    else:
        cnt = 0
        tmp_max_list = []
        for p in partitions:
            a, b = factorize(p)
            pool_shape = [1, a, b, 1]
            tmp_x = tf.reshape(x[:,cnt:cnt+p], pool_shape)
            tmp_max = tf.nn.max_pool(tmp_x, ksize=pool_shape, strides=pool_shape, padding='VALID')
            tmp_max_list.append(tmp_max)
            cnt += p
        tmp_max_len = len(tmp_max_list)
        assert tmp_max_len <= 15
        max_ = tf.concat(tmp_max_list, axis=2)
        max_ = tf.nn.max_pool(max_, ksize=[1,1,tmp_max_len,1], strides=[1,1,tmp_max_len,1], padding='VALID')
    x = tf.reshape(x, [-1, 1])
    x = tf.math.subtract(x, max_)

    x = tf.reshape(x, [1, -1])

    exp_x       = tf.exp(x)
    exp_sum     = tf.math.reduce_sum(exp_x, axis=-1)
    exp_sum_log = tf.log(exp_sum)

    return tf.math.subtract(x, exp_sum_log, name=name)

重要

训练用 tensorflow 的 log_softmax, 训练结束导出 CKPT 前改成 npu_log_softmax。

PyTorch

def factorize(n):
    for i in range(1, 16):
        if n % i == 0 and n // i <= 15:
            return (n // i, i)
    return ()

def split_and_factorize(n):
    result = []
    while not factorize(n):
        for i in range(n-1, 0, -1):
            if factorize(i):
                result.append(i)
                n -= i
                break
    result.append(n)
    return result

def npu_log_softmax(x, partitions):
    """ NPU LogSoftmax
    Args:
      x: A non-empty `Tensor`. limitation: len(x.shape) == 2 and x.shape[0] == 1
      partitions: A list of integers specifying the partition sizes for the log_softmax operation.
    Returns
      A `Tensor` representing the real log_softmax output.
    """

    if len(partitions) == 1:
        a, b       = factorize(partitions[0])
        pool_shape = [1, 1, a, b]
        x_   = x.reshape(pool_shape)
        max_ = torch.nn.functional.max_pool2d(x_, (a,b), (a,b))
    else:
        cnt = 0
        tmp_max_list = []

        for p in partitions:
            a, b = factorize(p)
            pool_shape = [1, 1, a, b]

            x_      = x[:,cnt:cnt+p]
            tmp_x   = x_.reshape(pool_shape)
            tmp_max = torch.nn.functional.max_pool2d(tmp_x, (a,b), (a,b))
            tmp_max_list.append(tmp_max)

            cnt += p

        tmp_max_len = len(tmp_max_list)

        max_ = torch.concat(tmp_max_list, axis=3)
        max_ = torch.nn.functional.max_pool2d(max_, (1,tmp_max_len), (1,tmp_max_len))

    max_  = max_.squeeze(0).squeeze(0)
    x     = x - max_

    exp_x       = torch.exp(x)
    exp_sum     = torch.sum(exp_x, dim=-1)
    exp_sum_log = torch.log(exp_sum)

    return x - exp_sum_log

重要

训练用 pytorch 的 log_softmax, 构建推理脚本时改成 npu_log_softmax。

4.3 BatchNorm*

可以设置 FUSE_BN 配置项选择是否把 BatchNorm 的参数合并到卷积

注意

BatchNorm 的参数合并到卷积后,如果开启卷积权重压缩,准确率可能会下降较多