TENSOR RT _ SIMPLE MNIST 예제

안녕하세요. 오늘은 TENSOR RT를 활용한 예시를 적어보려 합니다. TensorRT/samples/python/network_api_pytorch_mnist at master · NVIDIA/TensorRT (github.com) 공식 문서를 참고하였으며, 보기 좋게 Ipynb 를 사용하여 분석하였습니다.

 

 

IMPORT MODULE

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


 # TRT 모델을 작성하기 위한 모듈
import pycuda.autoinit
import tensorrt as trt
import pycuda.driver as cuda

 

 

BUILD MODEL 

- TRT 변환을 위한 모델을 만들어 준다. 원래 파이토치로 빌드 후, ONNX 파일을 빼내려 했지만, 제일 기초인 pytorch 학습을 완료한 후, weight 파일을 그대로 TRT 모델에 옮기는 작업을 진행해 보았다.

- 밑의 코드를 보게 되면 파이토치로 모델을 만들어준 다음, 학습이 완료한 Weight 를 Ordered Dict 형태로 model_weight 에 저장시켜 주었다. 이를 이용해서, TRT 모델에 weight를 집어 넣어 줄 것이다.

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.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.0025
        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,
            num_workers=1,
            timeout=600)
        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,
            num_workers=1,
            timeout=600)
        self.network = Net()

    # Train the network for one or more epochs, validating after each epoch.
    def learn(self, num_epochs=2):
        # Train the network for a single epoch
        def train(epoch):
            self.network.train()
            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.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()))

        # Test the network
        def test(epoch):
            self.network.eval()
            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()
                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
        
model = MnistModel()
model.learn()
weights = model.get_weights()

 

 

 

BUILD TRT LOGGER / BUILDER / NETWORK / RUNTIME 

- 이제 TRT 모델을 활용하기 위해서 Logger 와 builder 그리고, builder 를 통한 network를 정의해 준 후, runtime을 돌릴 것이다.

# Define Logger -> 가장 기본 세팅
TRT_Logger = trt.Logger(trt.Logger.WARNING)

EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

# Builder 정의 및 Network 설정
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network()
config = builder.create_builder_config()
runtime = trt.Runtime(TRT_LOGGER)

# config 의 GPU workspace 할당
config.max_workspace_size = 1 * 1 << 30  #setting gpu memory to 1GB

- 그 다음, 이제 가장 중요한 TRT Network 를 빌드해줄 것 인데, network 에 지원하는 layer 를 추가하여 만들어 주면 된다. input tensor 와 output tensor 를 설정해주었고, 이를 통해 나중에 input tensor 와 output tensor 를 바인딩 시켜주게 된다. 

- 핵심 부분만 보게 되면 다음과 같다.

# input tensor 정의
input_tensor = network.add_input(name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE)

# conv1.weight 추출
conv1_w = weights['conv1.weight'].numpy()
conv1_b = weights['conv1.bias'].numpy()

# TRT MODEL conv1 모델 빌드
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)

... 생략

 

 

 

BUILD ENGINE

- 만들어진 network 를 통해, Engine을 만들어주면 된다. Builder.build_engine( , ) 에는 2가지가 들어가게 되는데, network 와 config 를 넣어주면 된다. 이렇게 만들어진 engine 파일을 저장해주자.

# Build engine with serializing
engine_path = "./sample_mnist_api.engine"
with builder.build_engine(network,config) as engine,open(engine_path,"wb") as f:
    print(f"Serializing engine to file : {engine_path}")
    f.write(engine.serialize())

- 이 코드를 실행하면 .engine 의 파일이 저장되게 되는데, 이를 다시 deserialize 하게 되면, 모델 load 가 가능하게 된다.

with open(engine_path, "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())

 

 

 

PREPARE FOR INFERENCE

- 만들어진 engine은 단순 network 에 불과하다. 당연히, GPU 상에 모델이 올라가 있고, 이를 처리하기 위해서는 여러가지 처리 작업들이 필요하다. 단순 모델이다 보니, cuda 를 이용하여, CPU 상의 이미지를 GPU 상으로 옮겨 처리를 하게 하였다. 

- engine 에는 아까 설정해둔 , input 과 output 이 설정되어 있는데, 이를 cuda_memory_allocate를 통해, host (cpu) 와 device(GPU) 사이에서 메모리 할당을 해주어야 한다. 이부분은 Cuda 프로그래밍에 관한 부분이라 설명은 간단하게 하고 넘어가겠습니다.

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))
        # Allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size,dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        # Append the device buffer to 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

- 그 다음 inference 수행을 해주게 되면, 최종적으로 모델 output 이 나오게 됩니다.

def do_inference_v2(context, bindings, inputs, outputs, stream):
    # Transfer input data to the GPU.
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # Run inference.
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # Transfer predictions back from the GPU.
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # Synchronize the stream
    stream.synchronize()
    # Return only the host outputs.
    return [out.host for out in outputs]

- 최종 output 은 다음과 같이 같게 나오는 것을 확인할 수 있습니다.