[Model Inference] Torch-TensorRT 사용법 | 딥러닝 모델 최적화 및 인퍼런스 가속화

 

 

Torch-TensorRT는 PyTorch와 NVIDIA의 TensorRT를 통합하여 딥러닝 모델을 최적화하고 가속화하는 데 사용되는 PyTorch/TorchScript/FX용 컴파일러이다. Torch-TensorRT는 PyTorch 익스텐션으로 동작하며 JIT(Just In Time) 런타임에 원할하게 통합되는 모듈을 컴파일한다. 

 

NVIDIA TensorRT는 NVIDIA GPU에서 모델을 더 빠르게 실행하기 위한 최적화된 런타임 엔진으로, 특히 딥 러닝 모델을 배포 환경에서 더 효율적으로 실행하고 추론(inference) 성능을 향상시키는 데 사용된다.

 

기존에 파이썬으로 TensorRT를 사용하기 위해서는 오픈소스 커뮤니티에서 개발한 torch2trt 패키지를 사용해서 pytorch 모델을 tensorRT 호환 형식으로 변환해서 모델 인퍼런스를 가속화시켰다. 하지만 NIVIDA와 PyTorch가 공식적으로 제공하는 Torch-TensorRT를 사용하면  PyTorch 모델을 변환할 때 최적화 수준을 더 세밀하게 제어할 수 있다. 

Torch-TensorRT는 다음과 같은 주요 기능을 제공한다

• 모델 최적화: Torch-TensorRT는 PyTorch 모델을 TensorRT 호환 형식으로 변환하여 TensorRT의 최적화 기능을 활용할 수 있음. 이로써 모델의 연산 그래프가 최적화되고, GPU 가속을 위한 특정 연산 커널로 대체됨.
• TensorRT 모델 생성: Torch-TensorRT를 사용하여 최적화된 모델을 생성할 수 있음. 이렇게 생성된 모델은 TensorRT 엔진을 사용하여 효율적으로 추론을 수행할 수 있음.
• End-to-end 가속화: Torch-TensorRT를 사용하면 PyTorch 모델을 가져와 TensorRT 모델로 변환하고 추론을 수행하는 전체 엔드 투 엔드 파이프라인을 구축할 수 있음.
• 동적 그래프 지원: PyTorch의 동적 그래프 특성을 유지하면서 TensorRT로의 최적화가 가능하므로, 다양한 모델 아키텍처에 대한 지원을 제공합니다.

 

Torch-TensorRT를 사용하면 딥 러닝 모델을 NVIDIA GPU에서 더 효율적으로 실행하고, 추론 속도를 크게 향상시킬 수 있으며, 이는 실시간 요구 사항이 있는 응용 프로그램 및 서비스에 특히 유용하다. 

---

Torch-TensorRT Install

Torch-TensorRT — Torch-TensorRT v2.2.0.dev0+8ebf24d documentation

버전이 맞는 pytorch가 설치되어 있다면 tensorrt, torch_tensorrt 만 설치하면 예제 코드를 실행해 볼 수 있다.

pip install tensorrt torch_tensorrt

 

Torch-TensorRT 깃헙 Release 페이지

하지만 한번에 제대로 동작하지 않는 경우도 많다... ㅠ TensorRT는 환경 구성이 생각보다 힘들기 때문에 잘 동작하는 패키지 버전 조합을 잘 확인해서 설치해야 한다. Torch-TensorRT 깃헙 릴리스 페이지에서 Torch-TensorRT 버전별 pytorch, cuda, tensorrt와의 버전 호환성을 확인할 수 있다.

→ https://github.com/pytorch/TensorRT/releases

 

개인적으로는 TensorRT 8.6.1이 설치된 인스턴스를 띄워서 pytorch 2.0 cuda 11.8 을 설치한 뒤에 torch-tensorrt v1.4를 설치했다. 그랬더니... torchvision의 버전이 호환이 안된다고 해서 torchvision만 지웠다가 다시 설치하니까 잘 동작한다.

 

Python Torch-TensorRT 사용 예시

Torch-TensorRT로 입력을 컴파일하려면 torch.nn.ModuleTorch-TensorRT에 모듈과 입력을 제공하면 된다. 그러면 다른 PyTorch 모듈을 실행하거나 추가할 수 있도록 최적화된 TorchScript 모듈이 반환된다.

 

입력은 입력 텐서의 모양, 데이터 유형 및 메모리 형식을 정의하는 torch_tensorrt.Input 클래스 list이다. 입력이 튜플이나 Tensor 목록과 같은 더 복잡한 데이터 유형인 경우 인수를 사용하여 input_signature와 같은 컬렉션 기반 입력(e.g. List[Tensor], Tuple[Tensor, Tensor])을 지정할 수 있다(두 번째 예제 코드).

 

엔진이나 대상 장치의 작동 정밀도와 같은 설정을 지정할 수도 있고, 컴파일 후에는 다른 모듈과 마찬가지로 모듈을 저장하여 배포 애플리케이션에 로드할 수 있다. TensorRT/TorchScript 모듈을 로드하려면 먼저 torch_tensorrt를 임포트 해야 한다.

import torch_tensorrt

...

model = MyModel().eval()  # torch module needs to be in eval (not training) mode

inputs = [
    torch_tensorrt.Input(
        min_shape=[1, 1, 16, 16],
        opt_shape=[1, 1, 32, 32],
        max_shape=[1, 1, 64, 64],
        dtype=torch.half,
    )
]
enabled_precisions = {torch.float, torch.half}  # Run with fp16

trt_ts_module = torch_tensorrt.compile(
    model, inputs=inputs, enabled_precisions=enabled_precisions
)

input_data = input_data.to("cuda").half()
result = trt_ts_module(input_data)
torch.jit.save(trt_ts_module, "trt_ts_module.ts")

• 설치한 torch_tensorrt를 임포트
• pytorch 모델을 로드하고 torch_tensor.complie 수행
• 컴파일 시 input을 입력하는데, 이 때 모델 입력의 최소 크기,  최적화 크기, 최대 크기를 지정
• 이후 변환한 모델에 input을 입력해 주기만 하면 됨
• torch.jit.save로 변환한 모델 저장 가능

*입력 크기를 동적으로 조정할 수 있긴 하지만, 모델 정확도와 속도는 opt_shape (최적화된 크기)에서만 보장 받을 수 있음. 입력 크기가 opt_shape과 차이가 많이 날수록 pytorch 모델 출력과 tensorrt로 변환한 모델의 출력값의 차이가 많이 남.

 

# Sample using collection-based inputs via the input_signature argument
import torch_tensorrt

...

model = MyModel().eval()

# input_signature expects a tuple of individual input arguments to the module
# The module below, for example, would have a docstring of the form:
# def forward(self, input0: List[torch.Tensor], input1: Tuple[torch.Tensor, torch.Tensor])
input_signature = (
    [torch_tensorrt.Input(shape=[64, 64], dtype=torch.half), torch_tensorrt.Input(shape=[64, 64], dtype=torch.half)],
    (torch_tensorrt.Input(shape=[64, 64], dtype=torch.half), torch_tensorrt.Input(shape=[64, 64], dtype=torch.half)),
)
enabled_precisions = {torch.float, torch.half}

trt_ts_module = torch_tensorrt.compile(
    model, input_signature=input_signature, enabled_precisions=enabled_precisions
)

input_data = input_data.to("cuda").half()
result = trt_ts_module(input_data)
torch.jit.save(trt_ts_module, "trt_ts_module.ts")

• input_signature를 사용하는 예시
• 다양한 형태의 입력을 받을 수 있음

 

# Deployment application
import torch
import torch_tensorrt

trt_ts_module = torch.jit.load("trt_ts_module.ts")
input_data = input_data.to("cuda").half()
result = trt_ts_module(input_data)

• tensorRT로 변환 & 저장한 모델을 로드해서 사용하는 방법

 

# PyTorch 모델 추론
torch_output = model(input_data)

# TensorRT 모델 추론
trt_output = trt_ts_module(input_data)

# 두 출력 간의 유사성 확인
if torch.allclose(torch_output, trt_output, atol=1e-3):
    print("PyTorch와 TensorRT 모델의 출력이 유사합니다.")
else:
    print("PyTorch와 TensorRT 모델의 출력이 다릅니다.")

• pytorch 모델과 tensorrt 모델의 출력 유사성 테스트