Pytorch와 TensorRT를 이용한 딥 러닝 추론 최적화

이 글은 이전 글에서 계속되는 글입니다.

1. Pytorch Model

먼저 추론하기 위한 모델을 Pytorch를 이용하여 구성합니다. 앞 글에서 설명하였듯이 Bert를 이용합니다. Bert를 밑바닥부터 구성하긴 어려우니 transformers모듈에서 제공하는 모델을 사용합니다. 모델의 입력은 input_ids, attention_mask, token_type_ids 세 개를 받고, 출력은 Bert의 층 중 하나를 선택했습니다.
모델의 입력과 출력의 형태는 코드를 통해 확인할 수 있습니다.

import transformers
import torch

class Model(nn.Module):

    def __init__(self):
        super(Model, self).__init__()

        self.model = transformers.BertModel.from_pretrained('bert-base-cased')

    def forward(self, input_ids, attention_mask, token_type_ids):
        '''
        input : 
            input_ids : torch.as_tensor(x, dtype=torch.long), shape => (1, 512)
            attention_mask : torch.as_tensor(x, dtype=torch.long), shape => (1, 512)
            token_type_ids : torch.as_tensor(x, dtype=torch.long), shape => (1, 512)
        output :
            hiddens[-1] : torch.tensor(x, dtype=torch.float32), shape => (1, 512, 768)
        '''

        last_hiddens, last_pooling_hiddens, hiddens = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                token_type_ids=token_type_ids,
                output_hidden_states=True)

        return hiddens[-1]

2. Convert Pytorch to Onnx

Pytorch 모델을 Onnx로 변환하기 전 중요한 것은 Pytorch 모델에 input이 몇 개인지 output이 몇 개인지 확인하는 것입니다. 위에서 만든 Pytorch 모델은 3개의 input을, 1개의 output을 가지고 있습니다. 특이한 점은 실제 input과 같은 형태를 가지는 값을 넣어주어야 한다는 점입니다.
아래의 코드를 실행하면 bert.onnx라는 파일이 생성됩니다.

import numpy as np
import torch

model = Model().cuda().eval()
input_ids = torch.as_tensor(np.ones([1, 512]), dtype=torch.long).cuda()
attention_mask = torch.as_tensor(np.ones([1, 512]), dtype=torch.long).cuda()
token_type_ids = torch.as_tensor(np.ones([1, 512]), dtype=torch.long).cuda()

torch.onnx.export(
    model,
    (input_ids, attention_mask, token_type_ids),
    'bert.onnx',
    input_names=['input_ids', 'attention_mask', 'token_type_ids'],
    output_names=['outputs'],
    export_params=True)

3. Convert Onnx to TensorRT

변환된 Onnx파일을 TensorRT파일로 변환하여 저장합니다.
그리고 변환시 더 빠른 추론을 위해 float32로 되어 있는 부분을 float16으로 변환하는 설정도 넣어줍니다.
직접 해보시면 확인할 수 있지만 float32를 float16으로 변환하지 않으면 속도의 차이가 그렇게 크지 않습니다.

import tensorrt as trt

onnx_file_name = 'bert.onnx'
tensorrt_file_name = 'bert.plan'
fp_16_mode = True
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(EXPLICIT_BATCH)
parser = trt.OnnxParser(network, TRT_LOGGER)

builder.max_workspace_size = (1 << 30)
builder.fp16_mode = fp16_mode

with open(onnx_file_name, 'rb') as model:
    if not parser.parse(model.read()):
        for error in range(parser.num_errors):
            print (parser.get_error(error))

engine = builder.build_cuda_engine(network)
buf = engine.serialize()
with open(tensorrt_file_name, 'wb') as f:
    f.write(buf)

4. Load TensorRT File & TensorRT Inference

먼저 저장한 TensorRT 파일을 불러옵니다.

import tensorrt as trt
import pycuda.driver as cuda

tensorrt_file_name = 'bert.plan'
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt_runtime = trt.Runtime(TRT_LOGGER)

with open(tensorrt_file_name, 'rb') as f:
    engine_data = f.read()
engine = trt_runtime.deserialize_cuda_engine(engine_data)

불러온 TensorRT의 engine모듈을 분해하여 Stream으로 묶습니다. 그리고 각 텐서들을 GPU로 할당합니다.

class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()

inputs, outputs, bindings, 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_mem = cuda.pagelocked_empty*(size, dtype)
    device_mem = cuda.mem_alloc(host_mem.nbytes)
    bindings.append(int(device_mem))
    if engine.binding_is_input(binding):
        inputs.append(HostDeviceMem(host_mem, device_mem))
    else:
        outputs.append(HostDeviceMem(host_mem, device_mem))
context = engine.create_execution_context()

생성한 numpy array를 TensorRT가 읽을 수 있는 데이터로 변환 후 input의 입력으로 할당합니다.

input_ids, attention_mask, token_type_ids = np.ones([1, 512]), np.ones([1, 512]), np.ones([1, 512])

numpy_array_input = [input_ids, attention_mask, token_type_ids]
hosts = [input.host for input in inputs]
trt_types = [trt.int32, trt.int32, trt.int32]

for numpy_array, host, trt_types in zip(numpy_array_input, hosts, trt_types):
    numpy_array = np.asarray(numpy_array).astype(trt.nptype(trt_type)).ravel()
    np.copyto(host, numpy_array)

모든 데이터들을 streaming할 수 있도록 준비했으니 추론을 합니다.
여기서 TensorRT의 추론 출력은 numpy array의 shape가 맞지 않도록 나옵니다.
원하는 shape를 가지도록 정제하는 과정도 필요합니다.

def do_inference(context, bindings, inputs, outputs, stream):
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    stream.synchronize()
    return [out.host for out in outputs]

trt_outputs = do_inference(
                        context=context,
                        bindings=bindings,
                        inputs=inputs,
                        outputs=outputs,
                        stream=stream)
trt_outputs = trt_outputs.reshape((1, 512, 768))

5. Compare Inference time and embedding vector between Pytorch & TensorRT

Pytorch 모델과 TensorRT 모델을 이용하여 1000번씩 추론을 진행한 후 평균적으로 한번의 추론에 걸리는 시간을 비교합니다.
또한 실제로 Bert 모델이 출력하는 Embedding Vector를 비교하여 정확히 같은 값을 출력하도록 잘 변환이 되었는지 확인하였습니다.
물론 TensorRT로 변환하는 과정에서 float32를 모두 float16의 타입으로 변환하였기 때문에 정확히 같지는 않을 것 입니다.

pytorch time
0.018535836696624754
tensor([[[ 0.1215, -0.1506,  0.3299,  ..., -0.8052,  0.5446,  0.0536],
         [-0.0169, -0.2044,  0.5390,  ..., -0.0079,  0.7929, -0.1358],
         [-0.1034, -0.2198,  0.5454,  ..., -0.1563,  0.6997, -0.1943],
         ...,
         [-0.0149, -0.2066,  0.6120,  ..., -0.9506,  0.7732, -0.0641],
         [ 0.0634, -0.2646,  0.5229,  ..., -0.9125,  0.7374, -0.1188],
         [-0.0792, -0.2938,  0.5326,  ..., -0.8555,  0.6410, -0.2039]]],
       device='cuda:0', grad_fn=<AddcmulBackward>)
---------------------------------------------------------------------------
tensorrt time
0.007055652141571045
[[[ 0.11611275 -0.15233903  0.3321305  ... -0.80687803  0.5437621
    0.05408864]
  [-0.01713662 -0.20214725  0.53936094 ... -0.00627214  0.78893536
   -0.13889205]
  [-0.10317546 -0.21995904  0.5470193  ... -0.15895219  0.69504327
   -0.20108567]
  ...
  [-0.01862499 -0.21074559  0.61242217 ... -0.95111555  0.7742402
   -0.06390075]
  [ 0.06403332 -0.27156577  0.51832837 ... -0.91637695  0.7333
   -0.12652096]
  [-0.07960293 -0.2969381   0.5370592  ... -0.8596394   0.6355379
   -0.20504883]]]

위의 결과로 보아 매우 작은 숫자이지만 2배의 속도 향상이 있는 것으로 보여 속도 향상은 확실히 있는 것으로 알 수 있습니다.
실질적으로 Embedding Vector가 잘 추론되었는지 확인하기 위해 Vector의 Element들의 mean square error의 평균을 구했습니다.
1.208e-05의 mse로 나타나는 것을 확인할 수 있었으며 float32를 float16으로 변환한 것을 감안하면 잘 추론을 하는 것을 확인할 수 있습니다.

6. Reference

* https://news.developer.nvidia.com/tensorrt6-breaks-bert-record/

[##차금강##]