TensorFlow 笔记04

▶ 使用类封装写好的 TensorRT 模型，每个函数、类成员各司其职，而不是以前程序那样纯过程式，变量全部摊开

● 代码，程序入口 enter.py

import os
import sys
import numpy as np
import tensorrt as trt
import pycuda.autoinit
import pycuda.driver as cuda
from datetime import datetime as dt

import loadPara as ld
import calibrator

DEBUG           = True                     
testDataPath    = "./"
calibDataPath   = "./"
tempPath        = "./"
paraFile        = tempPath + "para.h5"
cacheFile       = tempPath + "calib.cache"
outputFile      = tempPath + "output.txt"
            
iGpu            = 0
calibCount      = 10                        # int8 校正次数
inputSize       = (1,1,1)                   # 输入数据尺寸，CHW

class TrtPredictor:
    def __init__(self, batchSize, dataType):
        self.logger     = trt.Logger(trt.Logger.ERROR)                      # 创建 logger
        self.batchSize  = batchSize
        self.dataType   = dataType
        self.h5f, ...   = fld.loadPara(paraFile)                            # 读取训练好的参数
        
        trtFilePath = tempPath + "engine-" + self.dataType + ".trt"         # 尝试读取创建好的引擎，没有则现场创建引擎
        if os.path.isfile(trtFilePath) and not DEBUG:
            f =  open(trtFilePath, 'rb')
            engineStr = f.read()                                            # enginStr 不作为成员变量
            self.runtime = trt.Runtime(self.logger)                         # 运行时读取文件中的引擎
            self.engine = self.runtime.deserialize_cuda_engine(engineStr)
            f.close()
            print("succeeded loading engine!")
        else:                             
            self.create_engine()                                            # 创建 engine，并写入文件，方便下次调用
            if self.engine == None:
                print("failed building engine!")
                return
            engineStr = self.engine.serialize()
            f = open(trtFilePath, 'wb')
            f.write(engineStr)
            f.close()
            print("succeeded building engine!")
                
        self.context = self.engine.create_execution_context()               # 创建 CUDA 上下文和流
        self.stream = cuda.Stream()
                
    def __del__(self):
        self.context = None
        self.engine  = None
        ld.close(self.h5f)
  
    def create_engine(self):                                                # 构造引擎
        self.builder = trt.Builder(self.logger)
        self.builder.max_batch_size     = 16
        self.builder.max_workspace_size = 1 << 30
        self.builder.fp16_mode          = self.dataType == 'float16'
        self.builder.int8_mode          = self.dataType == 'int8'
        self.network                    = self.builder.create_network()
        self.builder.strict_type_constraints = True        
               
        h0 = self.network.add_input("h0", trt.DataType.FLOAT, (1,) + inputSize) # 强制 N 为 1，多的数据堆在更高维度上
        
        #...                                                                # 中间层
        
        self.network.mark_output(h0.get_output(0))                          # 标记输出层

        if self.dataType == 'int8':                                         # int8 需要额外的校正，放到 builder 中
            self.builder.int8_calibrator = calibrator.MyCalibrator(calibCount, (self.batchSize,) + inputSize, calibDataPath, cacheFile)

        self.engine = self.builder.build_cuda_engine(self.network)          # 创建引擎（最容易失败的地方，返回构造函数后要检查是否成功）
    
    def infer(self, hInPart, dIn, dOut, hOut):                              # 推理
        cuda.memcpy_htod_async(dIn, hInPart, self.stream)
        self.context.execute_async(len(hInPart), [int(dIn), int(dOut)], self.stream.handle)
        cuda.memcpy_dtoh_async(hOut, dOut, self.stream)            
        self.stream.synchronize()
               
def predict(hIn, batchSize, dataType):    
    predictor = TrtPredictor(batchSize, dataType)                           # 构造一个预测器
    
    dIn  = cuda.mem_alloc(hIn[0].nbytes * batchSize)                        # 准备主机和设备内存
    hOut = np.empty((batchSize,) + tuple(predictor.engine.get_binding_shape(1)), dtype = np.float32)
    dOut = cuda.mem_alloc(hOut.nbytes)                                      # dOut 和 hOut 的大小一定是相同的
    res=[]
    for i in range(0, len(hIn), batchSize):                                 # 分 batch 喂入数据
        predictor.infer(hIn[i:i+batchSize], dIn, dOut, hOut)                    
        res.append( hOut )
    
    return res

if __name__ == "__main__":                                                  # main 函数负责管理 cuda.Device 和 cuda.Context
    _ = os.system("clear")
    batchSize = int(sys.argv[1])    if len(sys.argv) > 1 and sys.argv[1].isdigit()                         else 1
    dataType  = sys.argv[2]         if len(sys.argv) > 2 and sys.argv[2] in ['float32', 'float16', 'int8'] else 'float32'
    DEBUG     = int(sys.argv[3])>0  if len(sys.argv) > 3 and sys.argv[3].isdigit()                         else False
    if DEBUG:                                                               # 清除建好的 engine 和 校正缓存，重头开始建立                                
        oldEngineEAndCache = glob(tempPath+"*.trt") + glob(tempPath+"*.cache")
        [ os.remove(oldEngineEAndCache[i]) for i in range(len(oldEngineEAndCache))]
    print( "%s, start! GPU =  %s, batchSize = %2d, dataType  = %s" %( dt.now(), cuda.Device(iGpu).name(), batchSize, dataType ) )    
    
    inputData = loadData(testDataPath)                                      # 读取数据
    oF = open(outputFile, 'w')
    cuda.Device(iGpu).make_context()

    res = predict(inputData, batchSize, dataType)
    for i in range(len(res)):
        print( "%d -> %s" % (i,res[i]) )
        oF.write(res[i] + '
')
        
    oF.close()
    cuda.Context.pop()
    print( "%s, finish!" %(dt.now()) )

● 代码，矫正器 calibrator.py。核心思想是，手写一个数据生成器供 TensorRT 调用，每次从校正数据集中抽取 batchSize 那么多的数据，计算工作全部由 TensorRT 完成

import os
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibCount, inputShape, calibDataPath, cacheFile):
        trt.IInt8EntropyCalibrator2.__init__(self)                                              # 基类默认构造函数                                        
        self.calibCount     = calibCount                
        self.shape          = inputShape
        self.calibDataSet   = self.laodData(calibDataPath)                                      # 需要自己实现一个读数据的函数
        self.cacheFile      = cacheFile
        self.calibData      = np.zeros(self.shape, dtype=np.float32)        
        self.dIn            = cuda.mem_alloc(trt.volume(self.shape) * trt.float32.itemsize)     # 准备好校正用的设备内存      
        self.oneBatch       = self.batchGenerator()
            
    def batchGenerator(self):                                                                   # calibrator 的核心，一个提供数据的生成器
        for i in range(self.calibCount):
            print("> calibration ", i)
            self.calibData = np.random.choice(self.calibDataSet, self.shape[0], replace=False)  # 随机选取数据 
            yield np.ascontiguousarray(self.calibData, dtype=np.float32)                        # 调整数据格式后抛出   
    
    def get_batch_size(self):                                                                   # TensorRT 会调用，不能改函数名
        return self.shape[0]

    def get_batch(self, names):                                                                 # TensorRT 会调用，不能改函数名，老版本 TensorRT 的输入参数个数可能不一样
        try:
            data = next(self.oneBatch)                                                          # 生成下一组校正数据，拷贝到设备并返回设备地址，否则退出
            cuda.memcpy_htod(self.dIn, data)
            return [int(self.dIn)]
        except StopIteration:
            return None

    def read_calibration_cache(self):                                                           # TensorRT 会调用，不能改函数名
        if os.path.exists(self.cacheFile):
            print( "cahce file: %s" %(self.cacheFile) )
            f = open(self.cacheFile, "rb")
            cache = f.read()
            f.close()
            return cache              
            
    def write_calibration_cache(self, cache):                                                   # TensorRT 会调用，不能改函数名
        print( "cahce file: %s" %(self.cacheFile) )
        f = open(self.cacheFile, "wb")
        f.write(cache)
        f.close()

▶ 我的程序在 TensorRT 5 中 float32 和 float16 一切正常，int8 无法正确计算。具体表现为：正确加载 calibrator 调用，部分中间层计算结果与 float32 一模一样（二进制位层次上的相同，显然是采用了 float32 代替进行计算了），部分层所有计算结果与 float32 有分歧（10^-2 ~ 10^-3 量级上的），在之后多层计算中误差会逐渐放大，最终计算结果与 float32 大相径庭。更新 TensorRT 6 之后问题消失，int8 也能计算正确结果并获得加速。