Pytorch は手書き数字認識用の LeNet-5 ネットワークを構築し、tensorRT でデプロイします

序文

この記事は、私の学習モデル展開の簡単なデモです。単なる学習と記録のメモです。記事内の一部のコードとテキストはインターネットからのものです。侵害がある場合は、削除するために私に連絡してください。コード実装はすべて Python コードであり、C++ バージョンは実装されていません。

この記事では、環境のセットアップについては説明しません。コードの実行環境は次のとおりです。

pytorch=1.13
cuda=11.6
cudnn=8.8.0
tensorRT=8.5.3
pycuda=2022.2.2
opencv=4.7.0

これは私の以前の環境インストール記録ですcuda11.6.2 + cudnn8.8.0 + tensorRT8.5.3 + pytorch1.13 インストール記録 (プロテスト有効)

1. pytorch は LeNet-5 をビルドし、ONNX 形式に変換します

1.1 LeNet-5 ネットワークの紹介

論文「文書認識に適用された勾配ベース学習」の Lenet-5 ニューラル ネットワークは、手書き文字認識のための非常に効率的な畳み込みニューラル ネットワークです。Lenet-5 ニューラル ネットワークには合計 7 つの層があり、各層には異なる数のトレーニング パラメーターが含まれています。データのバッチをニューラル ネットワークに入力し、畳み込み、アクティブ化、プーリング、完全接続、ソフトマックス回帰などの後、最終的に確率配列を返し、画像を認識する目的を達成します。

1.2 ONNX (オープン ニューラル ネットワーク エクスチェンジ) の概要

Open Neural Network Exchange (Open Neural Network Exchange、ONNX) は、機械学習モデルを表すためのオープン標準ファイル形式であり、トレーニング済みモデルを保存するために使用できます。これにより、さまざまな機械学習フレームワーク (PyTorch、Caffe など) が可能になります。データは同じ形式で保存され、対話可能です。ONNX は、さまざまな機械学習モデルの対話性を強化するために、環境やプラットフォームに依存しない一連の標準形式を定義します。これにより、研究者は、あるフレームワークでモデルをトレーニングし、別のフレームワークで推論を行う自由が得られます。

ONNX の表現には、次の 2 つの主要な利点があります。

1. フレームワーク間の相互運用性

開発者は、さまざまなフレームワークをより簡単に切り替えて、さまざまなタスクに最適なツールを選択できるようになります。基本的に、各フレームワークは、トレーニング速度、ネットワーク アーキテクチャのサポート、モバイル デバイスでの推論など、特定の属性に合わせて最適化されています。ほとんどの場合、研究開発段階で最も必要とされる特性は、生産段階での特性とは異なります。これにより、最適なフレームワークに切り替えなかったり、モデルを別のフレームワークに移行して余分な作業が発生して進捗が遅れたりするなど、効率の低下につながります。ONNX 表現をサポートするフレームワークを使用すると、切り替えプロセスが大幅に簡素化され、開発者はより柔軟にツールを選択できるようになります。

2. 共有を最適化する

ハードウェア機器メーカーが開始したニューラル ネットワークのパフォーマンスの最適化は、ONNX を使用すると、複数の開発フレームワークに一度に影響を与えることができます。最適化が頻繁に行われる場合、最適化を各フレームワークに個別に統合するのは非常に時間がかかります。ONNX の表現を通じて、より多くの開発者がこれらの最適化にアクセスできるようになります。

1.3 pytorch は LeNet5 ネットワークを構築します

pytorch を使用して LeNet5 手書き数字認識を構築し、onnx 形式に変換します。

1. 関連モジュールをインポートする

import torch
import torch.nn as nn
import torch.optim as optim
import torch.onnx as onnx
import torch.nn.functional as F
import torchvision
from torch.utils.data import DataLoader

2. トレーニングパラメータを定義する

n_epochs = 3
batch_size_train = 64
batch_size_test = 1000
learning_rate = 0.01
momentum = 0.5
log_interval = 10
random_seed = 42
torch.manual_seed(random_seed)

3. データセットをロードする

train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.mnist.MNIST('./mnist/', train=True, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                               ])),
    batch_size=batch_size_train, shuffle=True)
test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.mnist.MNIST('./mnist/', train=False, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                               ])),
    batch_size=batch_size_test, shuffle=True)

4. ネットワーク構造を定義する

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)
    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)

5. 損失関数

# 定义模型和损失函数
network = Net()
optimizer = optim.SGD(network.parameters(), lr=learning_rate,
                      momentum=momentum)

train_losses = []
train_counter = []
test_losses = []
test_counter = [i*len(train_loader.dataset) for i in range(n_epochs + 1)]

6. トレーニングを開始する

def train(epoch):
    network.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = network(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))
            train_losses.append(loss.item())
            train_counter.append(
                (batch_idx * 64) + ((epoch - 1) * len(train_loader.dataset)))
            torch.save(network.state_dict(), './lenet5.pth')
            torch.save(optimizer.state_dict(), './optimizer.pth')


def test():
    network.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = network(data)
            test_loss += F.nll_loss(output, target, size_average=False).item()
            pred = output.data.max(1, keepdim=True)[1]
            correct += pred.eq(target.data.view_as(pred)).sum()
        test_loss /= len(test_loader.dataset)
        test_losses.append(test_loss)
        print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


for epoch in range(1, n_epochs + 1):
    train(epoch)
    test()

7. onnx形式に変換する

# 加载 PyTorch 模型
network.load_state_dict(torch.load("lenet5.pth"))

# 将 PyTorch 模型转为 ONNX 模型
input_shape = (1, 1, 28, 28)
dummy_input = torch.randn(input_shape)
onnx.export(network, dummy_input, "lenet5.onnx")

2. onnx を tensorRT に変換する

2.1 tensorRT の概要

TensorRTは、ディープ ラーニング アプリケーションに低遅延、高スループットのデプロイメント推論を提供できる、高性能のディープ ラーニング推論オプティマイザーおよびランタイム アクセラレーション ライブラリです。

TensorRT を使用すると、ハイパースケール データセンター、組み込みプラットフォーム、自動運転プラットフォームでの推論を高速化できます。

TensorRT は、TensorFlow、Caffe、Mxnet、Pytorch などのほぼすべての深層学習フレームワークをサポートできるようになり、TensorRT と NVIDIA GPU を組み合わせることで、ほぼすべてのフレームワークで高速かつ効率的なデプロイ推論を実行できます。

一般的なディープラーニングプロジェクトでは、トレーニングを高速化するために、マルチGPU分散トレーニングが使用されます。ただし、推論をデプロイする場合、コストを削減するために、単一の GPU マシンまたは組み込みプラットフォーム (NVIDIA Jetson など) をデプロイメントに使用することがよくあります。デプロイメント側にも、caffe などのトレーニングと同じ深層学習環境が必要です。 、TensorFlowなど。

トレーニングされたネットワーク モデルは大規模になる可能性があるため (開始、resnet など)、パラメーターが多く、デプロイメント側のマシンのパフォーマンスが異なるため、推論が遅くなり、遅延が大きくなります。これは、リアルタイムの目標検出や目標追跡などが必要な自動運転など、リアルタイム性の高いアプリケーションにとって致命的です。

導入推論の速度を向上させるために、モデル圧縮、枝刈り、量子化、知識蒸留などの多くのモデル最適化手法が登場しており、これらは一般にトレーニング段階で最適化されます。

TensorRT は、ネットワーク計算グラフを最適化することで、トレーニング済みモデルを最適化し、モデルの効率を向上させます。

ネットワークがトレーニングされた後、次のように、深層学習フレームワーク (Caffe、TensorFlow など) に依存せずに、トレーニング モデル ファイルを tensorRT に直接投入できます。

2.1 onnx から tensorRT への変換

import tensorrt as trt

TRT_LOGGER = trt.Logger()

# 加载 ONNX 模型
onnx_file_path = "lenet5.onnx"


def build_engine(onnx_file_path, engine_file_path):
    '''
    从ONNX文件创建TensorRT引擎以运行推理
    :return:
    '''
    # 创建一个TensorRT builder
    builder = trt.Builder(TRT_LOGGER)
    # 设置可由构建器使用的最大线程数
    builder.max_threads = 10
    flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    network = builder.create_network(flag)
    parser = trt.OnnxParser(network, TRT_LOGGER)
    runtime = trt.Runtime(TRT_LOGGER)
    # 设置可由运行时使用的最大线程数
    runtime.max_threads = 10

    # 解析模型文件
    with open(onnx_file_path, "rb") as model:
        print("开始ONNX文件解析")
        if not parser.parse(model.read()):
            print("错误：无法解析ONNX文件")
            for error in range(parser.num_errors):
                print(parser.get_error(error))

            return None

    print("完成ONNX文件解析")
    # 打印输入信息
    print("Network inputs:")
    for i in range(network.num_inputs):
        tensor = network.get_input(i)
        print(tensor.name, tensor.dtype, tensor.shape)

    config = builder.create_builder_config()
    config.set_flag(trt.BuilderFlag.REFIT)
    config.max_workspace_size = 1 << 28  # 256MiB

    plan = builder.build_serialized_network(network, config)
    engine = runtime.deserialize_cuda_engine(plan)

    with open(engine_file_path, "wb") as f:
        f.write(plan)

    return engine

build_engine(onnx_file_path, "lenet5.engine")

3. Opencv は画像をロードし、tensorRT を使用して推論を加速します

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2

TRT_LOGGER = trt.Logger()

# 加载 TensorRT Engine
engine_file_path = "lenet5.engine"
with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

# 创建执行上下文
context = engine.create_execution_context()

# 分配输入和输出内存
input_shape = (1, 1, 28, 28)
output_shape = (1, 10)
input_host = cuda.pagelocked_empty(trt.volume(input_shape), np.float32)
output_host = cuda.pagelocked_empty(trt.volume(output_shape), np.float32)
input_device = cuda.mem_alloc(input_host.nbytes)
output_device = cuda.mem_alloc(output_host.nbytes)

# 加载测试图像
img = cv2.imread(r"C:\Users\xr\Desktop\0.pgm", cv2.IMREAD_GRAYSCALE)
img = cv2.resize(img, (28, 28)) / 255.0

# 预处理输入图像
input_data = img.reshape((1,) + img.shape)
input_data = (input_data.astype(np.float32))
input_data = np.expand_dims(input_data, -1)
input_data = 1 - np.transpose(input_data, (0, 3, 1, 2))

# 将数据从主机内存复制到设备内存
np.copyto(input_host, input_data.ravel())
cuda.memcpy_htod(input_device, input_host)

# 执行 TensorRT Engine
context.execute_v2(bindings=[int(input_device), int(output_device)])

# 将数据从设备内存复制到主机内存
cuda.memcpy_dtoh(output_host, output_device)

# 后处理输出数据
prediction = np.argmax(output_host)
print(prediction)
# 显示结果
img = cv2.resize(img, (280, 280))
cv2.putText(img, "label:{}".format(prediction), (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0
cv2.imshow("input image", img)
cv2.waitKey()
cv2.destroyAllWindows()

推論結果