TensorRT&Sample&Python[network_api_pytorch_mnist]

---

本文是基於TensorRT 5.0.2基礎上，關於其內部的network_api_pytorch_mnist例子的分析和介紹。
本例子直接基於pytorch進行訓練，而後直接導出權重值爲字典，此時並未dump該權重；接着基於tensorrt的network進行手動設計網絡結構並填充權重。本文核心在於介紹network api的使用

1 引言

假設當前路徑爲：

TensorRT-5.0.2.6/samples

其對應當前例子文件目錄樹爲：

# tree python

python
├── common.py
├── network_api_pytorch_mnist
│   ├── model.py
│   ├── README.md
│   ├── requirements.txt
│   └── sample.py

2 基於pytorch

其中只有2個文件：

• model：該文件包含用於訓練Pytorch MNIST 模型的函數
• sample：該文件使用Pytorch生成的mnist模型去建立一個TensorRT inference engine

首先介紹下model.py

首先下載對應的mnist數據，並放到對應緩存路徑下：

'''
i) 去http://yann.lecun.com/exdb/mnist/index.html  下載四個
ii) 放到/tmp/mnist/data/MNIST/raw/
 '''

/tmp/mnist/data/MNIST/raw
├── t10k-images-idx3-ubyte.gz
├── t10k-labels-idx1-ubyte.gz
├── train-images-idx3-ubyte.gz
└── train-labels-idx1-ubyte.gz

這樣加快model.py讀取mnist數據的速度

# 該文件包含用於訓練Pytorch MNIST模型的函數
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable

import numpy as np
import os

from random import randint

# Network結構，2層卷積+dropout+一層全鏈接+一層softmax
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, kernel_size=5)
        self.conv2 = nn.Conv2d(20, 50, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(800, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.max_pool2d(self.conv1(x), kernel_size=2, stride=2)
        x = F.max_pool2d(self.conv2(x), kernel_size=2, stride=2)
        x = x.view(-1, 800)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)


class MnistModel(object):
    ''' 初始化'''
    def __init__(self):
        self.batch_size = 64
        self.test_batch_size = 100
        self.learning_rate = 0.01
        self.sgd_momentum = 0.9
        self.log_interval = 100

        # Fetch MNIST data set.
        # 訓練時候的數據讀取
        self.train_loader = torch.utils.data.DataLoader(
            datasets.MNIST('/tmp/mnist/data', train=True, download=True, transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,))
                ])),
            batch_size=self.batch_size,
            shuffle=True)

        # 測試時候的數據讀取
        self.test_loader = torch.utils.data.DataLoader(
            datasets.MNIST('/tmp/mnist/data', train=False, transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,))
                ])),
            batch_size=self.test_batch_size,
            shuffle=True)

        # 網絡結構實例化
        self.network = Net()


    ''' 訓練該網絡，而後每一個epoch以後進行驗證.'''
    def learn(self, num_epochs=5):

        # 每一個epoch的訓練過程
        def train(epoch):

            self.network.train()  # 開啓訓練flag
            optimizer = optim.SGD(self.network.parameters(), lr=self.learning_rate, momentum=self.sgd_momentum)

            for batch, (data, target) in enumerate(self.train_loader):
                data, target = Variable(data), Variable(target)
                optimizer.zero_grad()
                output = self.network(data)   # 一次前向
                loss = F.nll_loss(output, target)  # 計算loss
                loss.backward()  # 反向計算梯度
                optimizer.step()

                if batch % self.log_interval == 0:
                    print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                                            epoch, 
                                            batch * len(data), 
                                            len(self.train_loader.dataset), 
                                            100. * batch / len(self.train_loader), 
                                            loss.data.item()))

        # 測試該網絡
        def test(epoch):

            self.network.eval() # 開啓驗證flag
            test_loss = 0
            correct = 0

            for data, target in self.test_loader:
                with torch.no_grad():
                    data, target = Variable(data), Variable(target)
                output = self.network(data)  # 前向
                test_loss += F.nll_loss(output, target).data.item() # 累加loss值
                pred = output.data.max(1)[1]  # 計算當次預測值
                correct += pred.eq(target.data).cpu().sum() # 累加預測正確的

            test_loss /= len(self.test_loader)
            print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
                                             test_loss, 
                                             correct, 
                                             len(self.test_loader.dataset), 
                                            100. * correct / len(self.test_loader.dataset)))

        # 調用上面定義好的訓練函數和測試函數
        for e in range(num_epochs):
            train(e + 1)
            test(e + 1)

    ''' 可視化權重'''
    def get_weights(self):
        return self.network.state_dict()

    ''' 隨機獲取 測試樣本隊列中 樣本 '''
    def get_random_testcase(self):
        data, target = next(iter(self.test_loader))
        case_num = randint(0, len(data) - 1)
        test_case = data.numpy()[case_num].ravel().astype(np.float32)
        test_name = target.numpy()[case_num]
        return test_case, test_name

能夠看出，上面的代碼就是定義了網絡結構，和訓練網絡的函數方法。下面介紹下sample.py

# 該例子用pytorch編寫的MNIST模型去生成一個TensorRT Inference Engine
from PIL import Image
import numpy as np

import pycuda.driver as cuda
import pycuda.autoinit

import tensorrt as trt

import sys, os
sys.path.insert(1, os.path.join(sys.path[0], ".."))
import model


# import common
# 這裏將common中的GiB和find_sample_data，do_inference等函數移動到該py文件中，保證自包含。
def GiB(val):
    '''以GB爲單位，計算所須要的存儲值，向左位移10bit表示KB，20bit表示MB '''
    return val * 1 << 30

def find_sample_data(description="Runs a TensorRT Python sample", subfolder="", find_files=[]):
    '''該函數就是一個參數解析函數。
    Parses sample arguments.
    Args:
        description (str): Description of the sample.
        subfolder (str): The subfolder containing data relevant to this sample
        find_files (str): A list of filenames to find. Each filename will be replaced with an absolute path.
    Returns:
        str: Path of data directory.
    Raises:
        FileNotFoundError
    '''
    # 爲了簡潔，這裏直接將路徑硬編碼到代碼中。
    data_root = kDEFAULT_DATA_ROOT = os.path.abspath("/TensorRT-5.0.2.6/python/data/")

    subfolder_path = os.path.join(data_root, subfolder)
    if not os.path.exists(subfolder_path):
        print("WARNING: " + subfolder_path + " does not exist. Using " + data_root + " instead.")
    data_path = subfolder_path if os.path.exists(subfolder_path) else data_root

    if not (os.path.exists(data_path)):
        raise FileNotFoundError(data_path + " does not exist.")

    for index, f in enumerate(find_files):
        find_files[index] = os.path.abspath(os.path.join(data_path, f))
        if not os.path.exists(find_files[index]):
            raise FileNotFoundError(find_files[index] + " does not exist. ")

    if find_files:
        return data_path, find_files
    else:
        return data_path
#-----------------

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

class ModelData(object):
    INPUT_NAME = "data"
    INPUT_SHAPE = (1, 28, 28)
    OUTPUT_NAME = "prob"
    OUTPUT_SIZE = 10
    DTYPE = trt.float32


'''main中第三步：構建engine'''
# 該函數構建的網絡結構和上面model.py中一致，只是這裏經過訓練後的網絡模型讀取對應的權重值，並填充到network中
# network是TensorRT提供的，weights是Pytorch訓練後的模型提供的
def populate_network(network, weights):
 
    '''network支持的方法來自https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/infer/Graph/Network.html '''
    # 基於提供的權重配置網絡層
    input_tensor = network.add_input(name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE)

    conv1_w = weights['conv1.weight'].numpy()
    conv1_b = weights['conv1.bias'].numpy()
    conv1 = network.add_convolution(input=input_tensor, num_output_maps=20, kernel_shape=(5, 5), kernel=conv1_w, bias=conv1_b)
    conv1.stride = (1, 1)

    pool1 = network.add_pooling(input=conv1.get_output(0), type=trt.PoolingType.MAX, window_size=(2, 2))
    pool1.stride = (2, 2)

    conv2_w = weights['conv2.weight'].numpy()
    conv2_b = weights['conv2.bias'].numpy()
    conv2 = network.add_convolution(pool1.get_output(0), 50, (5, 5), conv2_w, conv2_b)
    conv2.stride = (1, 1)

    pool2 = network.add_pooling(conv2.get_output(0), trt.PoolingType.MAX, (2, 2))
    pool2.stride = (2, 2)

    fc1_w = weights['fc1.weight'].numpy()
    fc1_b = weights['fc1.bias'].numpy()
    fc1 = network.add_fully_connected(input=pool2.get_output(0), num_outputs=500, kernel=fc1_w, bias=fc1_b)

    relu1 = network.add_activation(input=fc1.get_output(0), type=trt.ActivationType.RELU)

    fc2_w = weights['fc2.weight'].numpy()
    fc2_b = weights['fc2.bias'].numpy()
    fc2 = network.add_fully_connected(relu1.get_output(0), ModelData.OUTPUT_SIZE, fc2_w, fc2_b)

    fc2.get_output(0).name = ModelData.OUTPUT_NAME
    network.mark_output(tensor=fc2.get_output(0))


'''main中第三步：構建engine'''
def build_engine(weights):

    '''下面的create_network會返回一個tensorrt.tensorrt.INetworkDefinition對象
     https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/infer/Core/Builder.html?highlight=create_network#tensorrt.Builder.create_network 
    '''

    with trt.Builder(TRT_LOGGER) as builder, \
           builder.create_network() as network:

        builder.max_workspace_size = GiB(1)

        populate_network(network, weights)   # 用以前的pytorch模型中的權重來填充network

        # 構建並返回一個engine.
        return builder.build_cuda_engine(network)


'''main中第四步：分配buffer '''
def allocate_buffers(engine):

    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()

    for binding in engine:

        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))

        # 分配host和device端的buffer
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)

        # 將device端的buffer追加到device的bindings.
        bindings.append(int(device_mem))

        # Append to the appropriate list.
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))

    return inputs, outputs, bindings, stream


'''main中第五步：選擇測試樣本 '''
# 用pytorch的DataLoader隨機選擇一個測試樣本
def load_random_test_case(model, pagelocked_buffer):

    img, expected_output = model.get_random_testcase()

    # 將圖片copy到host端的pagelocked buffer
    np.copyto(pagelocked_buffer, img)

    return expected_output


'''main中第六步：執行inference '''
# 該函數能夠適應多個輸入/輸出；輸入和輸出格式爲HostDeviceMem對象組成的列表
def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):

    # 將數據移動到GPU
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]

    # 執行inference.
    context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)

    # 將結果從 GPU寫回到host端
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]

    # 同步stream
    stream.synchronize()

    # 返回host端的輸出結果
    return [out.host for out in outputs]


def main():

    ''' 1 - 尋找模型文件，不過次例中未用到該返回值'''
    data_path = find_sample_data(description="Runs an MNIST network using a PyTorch model", subfolder="mnist")

    ''' 2 - 訓練該模型'''
    mnist_model = model.MnistModel() 
    mnist_model.learn()

    # 獲取訓練好的權重
    weights = mnist_model.get_weights()

    ''' 3 - 基於build_engine構建engine；用tensorrt來進行inference '''
    with build_engine(weights) as engine:

        ''' 4 - 構建engine, 分配buffers， 建立一個流 '''
        inputs, outputs, bindings, stream = allocate_buffers(engine)

        with engine.create_execution_context() as context:

            ''' 5 - 讀取測試樣本，並歸一化'''
            case_num = load_random_test_case(mnist_model, pagelocked_buffer=inputs[0].host)

            ''' 6 -執行inference，do_inference函數會返回一個list類型，此處只有一個元素 '''
            [output] = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
            pred = np.argmax(output)

            print("Test Case: " + str(case_num))
            print("Prediction: " + str(pred))

if __name__ == '__main__':
    main()

運行結果以下：