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 格式
-
参数格式
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.yamlINPUT_OPS: Feats: [1, 1, 64] State_c0: [1, 3, 64] State_c1: [1, 4, 64] State_c2: [1, 5, 64]
-
参数格式
[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.yamlOUTPUT_OPS: [State_c0_out, State_c1_out, State_c2_out, Result]
注意事项
- 输出状态结点必须放在预测输出节点之前
-
参数相关概述
NPU 内部运算采用 FP16 格式数据,输入输出张量都是 FP16 格式;
NPU 支持格式转换 FP16_TO_FP32 功能,不支持格式转换 FP32_TO_FP16 功能;
循环神经网络的处理流程中,在每个时间步下,网络会接收当前时刻的输入以及前一时间步的隐藏状态,通过这些信息生成当前时刻的隐藏状态和相应的预测结果;
在实际应用中,预测结果对应的输出张量会先转换为 FP32 格式后,再进行预测处理,隐藏状态对应的输出张量则直接以 FP16 格式作为下一帧的输入;
用户需要按照具体模型结构,将输出节点的数据格式划分为 FP32 和 FP16;
用户可通过 FP16_OUT_OPS 参数指定哪些输出节点作为 FP16 格式输出;
-
参数格式
[output_state_name, ...]
output_state_name 为模型隐藏状态对应的输出节点名称
-
使用示例
该模型配置文件中的 FP16_OUT_OPS 参数可按如下方式配置:
config.yamlFP16_OUT_OPS: [State_c0_out, State_c1_out, State_c2_out]
-
参数相关概述
NPU 编译器权重量化功能开启后,若模型推理性能较差,用户可根据权重分布直方图,分析各权重节点是否量化友好,并可指定哪些权重节点不进行量化处理
-
参数格式
[weight_op_name, ...]
weight_op_name 为需要关闭量化处理的权重节点名称
-
使用示例
本例中,假设模型中部分卷积权重量化处理后,导致整体模型性能较差,故不对其进行量化处理
配置文件中的 EXCLUDE_COMPRESS_OPS 参数可按如下方式配置:
config.yamlEXCLUDE_COMPRESS_OPS: [conv2d_5/Conv2D/ReadVariableOp/_74__cf__74, conv2d_6/Conv2D/ReadVariableOp/_75__cf__75]
-
参数相关概述
NPU 编译器采用训练后量化(PTQ, Post Train Quantization),量化策略选择 MinMax 方法统计权重数据分布范围;
NPU 编译器权重量化功能开启后,若模型推理性能较差,用户可根据权重分布直方图,分析各权重节点数据分布范围的统计是否合理,并可通过该参数直接配置指定权重节点的数据分布范围;
-
参数格式
weight_op_name: [min, max]
weight_op_name —— 模型指定某个权重节点的名称 min, max —— 指定 weight_op_name 张量的最大最小值
-
使用示例
本例中指定两组卷积的量化处理范围,配置文件中的 EXCLUDE_COMPRESS_OPS 参数可按如下方式配置:
config.yamlWEIGHT_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_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.yamlINPUT_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)
-
参数相关概述
在深度学习领域中,模型算子节点之间通常采用多维度张量进行数据传输,比如卷积神经网络的特征图通常采用四维度张量存储
四维度张量的各维度可分别表示为 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.yamlINPUT_NCX_TO_NXC: [0] INPUT_OPS: 0: [1, 32, 128] 1: [1, 1, 128]
下图为模型结构图:
-
不使能参数,输入张量 input_tensor 保持原始数据排布格式
模型配置文件的 INPUT_NCX_TO_NXC 和 INPUT_OPS 参数可按如下方式配置:
config.yamlINPUT_NCX_TO_NXC: [] INPUT_OPS: 0: [1, 128, 32] 1: [1, 1, 128]
下图为模型结构图:
对比上述两 NPU 模型结构图可知,格式转换功能开启后,可将输入侧插入的转置节点优化
该参数能够根据用户需求,更加灵活的设置输入张量的数据格式
-
-
参数相关概述
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.yamlFP16_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 的参数合并到卷积后,如果开启卷积权重压缩,准确率可能会下降较多