TensorRT (tensor runtime) và ứng dụng

Khái niệm

TensorRT (Tensor Runtime) là một thư viện tối ưu hóa mô hình deep learning phát triển bởi NVIDIA. Nó cung cấp một loạt các công cụ để tối ưu hóa và triển khai mô hình deep learning trên các thiết bị có GPU của NVIDIA, nhằm cải thiện hiệu suất và giảm độ trễ khi chạy mô hình trên các ứng dụng thời gian thực.

Khi sử dụng với PyTorch, TensorRT có thể được sử dụng để tối ưu hóa mô hình PyTorch trước khi triển khai chúng. Quá trình này thường bao gồm các bước sau:

1. Chuyển đổi mô hình PyTorch thành mô hình TorchScript hoặc ONNX: TensorRT hỗ trợ các định dạng đầu vào này, vì vậy trước tiên, bạn cần chuyển đổi mô hình PyTorch của mình thành một trong các định dạng này.

2. Tối ưu hóa mô hình bằng TensorRT: Sau khi có mô hình TorchScript hoặc ONNX, bạn có thể sử dụng TensorRT để tối ưu hóa mô hình. TensorRT sẽ thực hiện một loạt các biến đổi nhằm cải thiện hiệu suất của mô hình trên GPU.

3. Triển khai mô hình đã được tối ưu hóa: Mô hình đã được tối ưu hóa có thể được triển khai trên các ứng dụng hoạt động thời gian thực hoặc được tích hợp vào hệ thống nhúng.

Việc sử dụng TensorRT cung cấp nhiều lợi ích, bao gồm:

• Tăng tốc độ chạy mô hình: Mô hình tối ưu hóa bằng TensorRT thường chạy nhanh hơn so với mô hình gốc trên GPU của NVIDIA.

• Giảm độ trễ: TensorRT tối ưu hóa mô hình để giảm độ trễ khi chạy, điều này đặc biệt quan trọng đối với các ứng dụng thời gian thực.

• Hỗ trợ nhiều loại mô hình: TensorRT hỗ trợ nhiều kiến trúc mạng khác nhau, bao gồm CNNs, RNNs và mạng kết hợp.

Tóm lại, TensorRT là một công cụ hữu ích để cải thiện hiệu suất và độ trễ của mô hình deep learning trên GPU của NVIDIA, và khi sử dụng với PyTorch, nó cung cấp một cách tiếp cận tối ưu hóa mô hình hiệu quả.

Ví dụ

Để sử dụng TensorRT để tối ưu hóa mô hình image classification được xây dựng bằng PyTorch, bạn cần thực hiện các bước sau:

1. Huấn luyện và lưu mô hình PyTorch:

   Trước tiên, bạn cần huấn luyện mô hình image classification bằng PyTorch. Sau đó, lưu mô hình dưới dạng TorchScript hoặc ONNX. Đây là ví dụ về cách lưu mô hình dưới dạng TorchScript:

   1
   2
   3
   4
   5
   6
   7
   8
   9
   10
   import torch
   import torchvision.models as models
      
   # Tạo và huấn luyện mô hình
   model = models.resnet18(pretrained=True)
   example_input = torch.randn(1, 3, 224, 224)  # Example input shape (batch_size, channels, height, width)
   traced_script_module = torch.jit.trace(model, example_input)
      
   # Lưu mô hình dưới dạng TorchScript
   traced_script_module.save("resnet18.pt")
   

2. Sử dụng TensorRT để tối ưu hóa mô hình:

   Sau khi có mô hình TorchScript hoặc ONNX, bạn có thể sử dụng TensorRT để tối ưu hóa mô hình. Dưới đây là ví dụ về cách sử dụng TensorRT với TorchScript:

   1
   2
   3
   4
   5
   6
   7
   8
   9
   10
   11
   12
   import torch
   from torch2trt import torch2trt
   from torch2trt.module_test import add_module_test
      
   # Load mô hình TorchScript đã lưu
   model = torch.jit.load("resnet18.pt")
      
   # Tối ưu hóa mô hình bằng TensorRT
   model_trt = torch2trt(model, [torch.randn(1, 3, 224, 224).cuda()])
      
   # Lưu mô hình đã tối ưu hóa
   torch.save(model_trt.state_dict(), "resnet18_trt.pth")
   

   Trong đoạn mã trên, torch2trt được sử dụng để chuyển đổi mô hình TorchScript thành mô hình đã được tối ưu hóa bằng TensorRT.

3. Triển khai mô hình đã tối ưu hóa:

   Mô hình đã được tối ưu hóa có thể được triển khai trong ứng dụng của bạn. Bạn có thể sử dụng torch2trt để tối ưu hóa mô hình và sau đó sử dụng mô hình đã tối ưu hóa trong ứng dụng của mình.

Nhớ cài đặt các thư viện cần thiết trước khi thực hiện các bước trên. Ngoài ra, hãy kiểm tra phiên bản của PyTorch và TensorRT để đảm bảo tích hợp tốt.

Load Model

Để load và sử dụng mô hình đã được tối ưu hóa bằng TensorRT từ tệp resnet18_trt.pth, bạn cần thực hiện các bước sau:

1. Import các thư viện cần thiết:

1
2
3
import torch
import torchvision.transforms as transforms
from PIL import Image

1. Định nghĩa hàm để load và sử dụng mô hình:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def load_trt_model(trt_model_path):
    # Tạo một mô hình ResNet18 trống
    model = torchvision.models.resnet18(pretrained=False)
    # Load trạng thái đã tối ưu hóa từ tệp resnet18_trt.pth
    model.load_state_dict(torch.load(trt_model_path))
    # Chuyển mô hình vào chế độ đánh giá (evaluation mode)
    model.eval()
    return model

def preprocess_image(image_path):
    # Chuẩn bị ảnh đầu vào cho mô hình
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    image = Image.open(image_path)
    image = preprocess(image).unsqueeze(0)
    return image

def classify_image(model, image_path, class_names):
    image = preprocess_image(image_path)
    with torch.no_grad():
        outputs = model(image)
    _, predicted = outputs.max(1)
    class_name = class_names[predicted.item()]
    return class_name

if __name__ == "__main__":
    # Đường dẫn tới tệp resnet18_trt.pth
    trt_model_path = "resnet18_trt.pth"

    # Load mô hình đã được tối ưu hóa
    trt_model = load_trt_model(trt_model_path)

    # Danh sách tên lớp
    class_names = ["class1", "class2", "class3", ...]  # Thay thế bằng tên lớp thực tế

    # Đường dẫn tới ảnh cần phân loại
    image_path = "your_image.jpg"

    # Phân loại ảnh
    result = classify_image(trt_model, image_path, class_names)
    print(f"Predicted class: {result}")

Lưu ý rằng bạn cần thay thế class_names bằng danh sách tên lớp thực tế mà mô hình của bạn đã được huấn luyện để phân loại. Điều này giúp bạn hiển thị kết quả phân loại dưới dạng tên lớp thay vì chỉ số.

Tài liệu tham khảo

Internet

Hết.

• Previous
  Review satellite communications in the new space era: a survey and future challenges
• Next
  How to custom layer from architecture model in pytorch