Dubhe/model_compress/docs/API_quant.md

• OneFlow 量化推理
• • OneFlow 中的 XRT
  • 在 OneFlow 中使用 TensorRT
  • 环境
  • • 单机单卡

OneFlow 量化推理

OneFlow 中的 XRT

XRT 是一个同时支持多个计算引擎的运行时加速库，目前已经集成了 TensorFlow XLA 和 NVIDIA
TensorRT 两个后端引擎。其中 XLA 全面支持训练和预测，TensorRT 支持预测以及部分算子支持训练。对于同一个计算图，XRT 允许多个计算引擎联合使用，以获得更好的加速效果，其中 TensorRT 具有 Int8 量化功能。

由于 TensorRT 中官方支持的 op 并没有那么全面，其余自定义 op 有可能受接口限制，因此 OneFlow 后续会采用 plug-in 形式添加，支持更多算子。

在 OneFlow 中使用 TensorRT

• 前期准备
  • 数据集：以测试 ResNet50 为例，需要提前准备 ImageNet 的 OFRecord 格式数据集。
  • 下载 TensorRT：编译时需要链接 TensorRT 的头文件和动态库，因此用户需要根据自己系统和已安装的 CUDA 版本选择相应版本的 TensorRT，同时满足 TensorRT 的其他依赖。
  • 下载 OneFlow-Benchmark：OneFlow-Benchmark 是 OneFlow 的模型基准仓库，提供了一系列完备实现的网络模型，本次测试选择的是其中的ResNet50。
• 编译：编译时开启 -DWITH_TENSORRT 选项，并指定 TensorRT 源码解压后的所在路径

cmake .. -DWITH_TENSORRT=ON -DTENSORRT_ROOT=/home/${user}/TensorRT-6.0.1.8 && make -j 24

或者可以在 cmake 前使用环境变量指定

export TENSORRT_ROOT=/home/${user}/TensorRT-6.0.1.8

编译成功后即可安装支持 TensoRT 的 OneFlow。

• 运行
  目前 OneFlow 中的 TensorRT 仅支持单卡推理。编译成功后切换到 dev_trt_infer 分支，在 config.py 中
• 添加 --use_tensorrt，可使用 TenosrRT 推理。
• 添加 --use_tensorrt 和 use_int8，可开启 TenosrRT 的 int8 量化。

环境

硬件环境

• CPU：Intel(R) Xeon(R) CPU E5-2650 v4 @ 2.20GHz x 6
• GPU：[GeForce GTX 1080] x 4

软件环境

• 系统：Ubuntu 18.04.4 LTS
• NVIDIA Driver Version：440.44
• CUDA：10.2
• GCC：7.5
• Cmake：3.14.4
• Make：4.1

测试结果

测试模型为 ResNet 50（以下称 rn50），使用在线量化，分别进行单机单卡和单机多卡推理，batch_size 取 64 和可运行的最大 batch_size。
若正常运行，log 打印如下：

==================================================================
Running resnet50: num_gpu_per_node = 1, num_nodes = 1.
==================================================================
dtype = float32
gpu_num_per_node = 1
num_nodes = 1
node_ips = ['127.0.0.1']
ctrl_port = 50051
model = resnet50
use_fp16 = None
use_xla = None
channel_last = None
pad_output = None
num_epochs = 1
model_load_dir = resnet_v15_of_best_model_val_top1_77318
batch_size_per_device = 64
val_batch_size_per_device = 256
nccl_fusion_threshold_mb = 0
nccl_fusion_max_ops = 0
fuse_bn_relu = False
fuse_bn_add_relu = False
gpu_image_decoder = False
image_path = test_img/tiger.jpg
num_classes = 1000
num_examples = 1281167
num_val_examples = 50000
rgb_mean = [123.68, 116.779, 103.939]
rgb_std = [58.393, 57.12, 57.375]
image_shape = [3, 224, 224]
label_smoothing = 0.1
model_save_dir = ./output/snapshots/model_save-20201123172206
log_dir = ./output
loss_print_every_n_iter = 1
image_size = 224
resize_shorter = 256
train_data_dir = None
train_data_part_num = 256
val_data_dir = /dataset/ImageNet/ofrecord/validation
val_data_part_num = 256
optimizer = sgd
learning_rate = 0.256
wd = 3.0517578125e-05
momentum = 0.875
lr_decay = cosine
lr_decay_rate = 0.94
lr_decay_epochs = 2
warmup_epochs = 5
decay_rate = 0.9
epsilon = 1.0
gradient_clipping = 0.0
------------------------------------------------------------------
Time stamp: 2020-11-23-17:22:06
Restoring model from resnet_v15_of_best_model_val_top1_77318.
Loading data from /dataset/ImageNet/ofrecord/validation


W1123 17:23:41.120939 31217 trt_executable.cpp:146] Rebuild engine since the maximum batch size 1 is less than the input batch size 256
W1123 17:24:25.756124 33076 trt_logger.cpp:35] TensorRT Logging: Explicit batch network detected and batch size specified, use execute without batch size instead.
W1123 17:24:31.005220 33076 trt_logger.cpp:35] TensorRT Logging: Explicit batch network detected and batch size specified, use execute without batch size instead.
W1123 17:24:36.085610 33076 trt_logger.cpp:35] TensorRT Logging: Explicit batch network detected and batch size specified, use execute without batch size instead.
W1123 17:24:41.073289 33076 trt_logger.cpp:35] TensorRT Logging: Explicit batch network detected and batch size specified, use execute without batch size instead.
W1123 17:24:45.920917 33076 trt_logger.cpp:35] TensorRT Logging: Explicit batch network detected and batch size specified, use execute without batch size instead.
W1123 17:24:50.633805 33076 trt_logger.cpp:35] TensorRT Logging: Explicit batch network detected and batch size specified, use execute without batch size instead.
W1123 17:24:55.354147 33076 trt_logger.cpp:35] TensorRT Logging: Explicit batch network detected and batch size specified, use execute without batch size instead.
W1123 17:24:59.904863 33076 trt_logger.cpp:35] TensorRT Logging: Explicit batch network detected and batch size specified, use execute without batch size instead.
validation: epoch 0, iter 195, top_1: 0.772155, top_k: 0.934856, samples/s: 181.038 1606123666.3968866

单机单卡