Die Methode und Anwendung zum Konvertieren des PyTorch-Modells in das ONNX-Modell

Dieser Artikel stellt am Beispiel der Ziffernerkennung des MNIST-Bilddatensatzes den grundlegenden Prozess des Trainings des CNN-Modells mit dem PyTorch-Framework, die Methode zur Konvertierung des PyTorch-Modells in das ONNX-Modell und die Funktionsweise des ONNX-Modells vor.

Import des MNIST-Datensatzes

MNIST-Datensatz-URL:
MNIST-Datensatz

Der Code zum Importieren des Datensatzes lautet wie folgt:

import torch
import torch.utils.data as data
import torchvision
import matplotlib.pyplot as plt


DOWNLOAD_MNIST = False  # if need to download data, set True at first time


# read train data
train_data = torchvision.datasets.MNIST(
    root='./data', train=True, download=DOWNLOAD_MNIST, transform=torchvision.transforms.ToTensor())
print()
print("size of train_data.train_data:  {}".format(train_data.train_data.size()))  # train_data.train_data is a Tensor
print("size of train_data.train_labels:  {}".format(train_data.train_labels.size()), '\n')

# plot one example
plt.imshow(train_data.train_data[50].numpy(), cmap='Greys')
plt.title('{}'.format(train_data.train_labels[50]))
plt.show()

# data loader
# combines a dataset and a sampler, and provides an iterable over the given dataset
train_loader = data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)

# read test data
test_data = torchvision.datasets.MNIST(root='./data', train=False)
num_test, num_test_over = 2000, 1000
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(
    torch.FloatTensor)[: num_test] / 255.  # unsqueeze because of 1 channel; value in range (0, 1)
test_y = test_data.test_labels[: num_test]
test_over_x = torch.unsqueeze(
    test_data.test_data, dim=1).type(torch.FloatTensor)[-num_test_over:] / 255.  # test data after training
test_over_y = test_data.test_labels[-num_test_over:]

Vorsichtsmaßnahmen:

1. Der Download-Parameter von Torchvision.datasets.MNIST wird auf True gesetzt, wenn der Datensatz heruntergeladen werden muss, und auf False gesetzt, wenn der Datensatz heruntergeladen wurde;

Achten Sie zweitens auf den Dateipfad des Datensatzes:

Drittens kann Torchvision.datasets.MNIST sowohl Trainingsdaten als auch Testdaten lesen. Beim Lesen von Trainingsdaten sind die train_data des erstellten Objekts Tensor-Daten, und die Objekte werden durch den Torch.utils.data.DataLoader-Sampling-Vorgang gestapelt; beim Lesen von Test data, der train-Parameter wird während der Instanziierung auf False gesetzt, und die test_data des erstellten Objekts sind Tensor-Daten.

Training und Testen von CNN-Modellen

Dieser Artikel verwendet ein relativ einfaches CNN-Modell, das darauf trainiert ist, Ziffern im MNIST-Datensatz zu erkennen.

Der PyTorch-Modellcode lautet wie folgt:

import torch.nn as nn


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # for 28*28 to 14*14
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # for 14*14 to 7*7
        )
        self.output = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = out.view(out.size(0), -1)  # flatten the 32*7*7 dimension to 1568
        out = self.output(out)

        return out

Darunter sind die MNIST-Daten alle Bilder der Größe 28 x 28 nach zwei Faltungsmodulen (jedes Modul durchläuft eine Faltungsschicht, eine nichtlineare ReLU-Schicht und eine Max-Pooling-Schicht) und geben einen Vektor mit einer Länge von 10 aus, was das Gewicht von 10 Zahlen.

Der Code des Modelltrainingsteils lautet wie folgt:

import torch.nn as nn

from model import CNN


cnn = CNN()
print("CNN model structure:\n")
print(cnn, '\n')
optimizer = torch.optim.Adam(params=cnn.parameters(), lr=LR)  # optimizer: Adam
loss_func = nn.CrossEntropyLoss()  # loss function: cross entropy

for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):  # x, y: Tensor of input and output
        output = cnn(x)
        loss = loss_func(output, y)

        # optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # test at each 50 steps
        if not step % 50:
            test_output = cnn(test_x)
            predict_y = torch.max(test_output, 1)[1].data.numpy()
            accuracy = float((predict_y == test_y.data.numpy()).astype(int).sum()) / float(test_y.size(0))
            print('Epoch: ', epoch, '| train loss: %.4f' % loss.data.numpy(), '| test accuracy: %.3f' % accuracy)


# output model file
torch.save(cnn, 'cnn.pt')
print()
print('finish training', '\n')

Nach dem Training kann das Modell in der lokalen .pt-Datei über Torch.save gespeichert werden, bei der es sich um die PyTorch-Modelldatei handelt.Wie die Modelldatei in das ONNX-Modell konvertiert wird, wird später beschrieben.

Der Code des Modelltestteils lautet wie folgt:

# load model
print('load cnn model', '\n')
cnn_ = torch.load('cnn.pt')

# test new data
test_output = cnn_(test_over_x)
predict_y = torch.max(test_output, 1)[1].data.numpy()
print("prediction number:  {}".format(predict_y))
print("real number:  {}".format(test_over_y.numpy()), '\n')
accuracy = float((predict_y == test_over_y.data.numpy()).astype(int).sum()) / float(test_over_y.size(0))
print("accuracy:  {}".format(accuracy), '\n')

Konvertierung des PyTorch-Modells in das ONNX-Modell und die Anwendung

Der vollständige Name von ONNX lautet Open Neural Network eXchange, das eine offene Standardform des maschinellen Lernmodells ist, die Modellkonvertierung mehrerer Frameworks unterstützt und eine große Vielseitigkeit aufweist.
Darüber hinaus ist die Betriebseffizienz (Inferenz) des ONNX-Modells deutlich besser als die des PyTorch-Modells.

Code zeigen wie folgt:

import os

import numpy as np
import torch
import torchvision
import onnx
import onnxruntime as ort

from model import CNN  # need model structure


PATH = os.path.dirname(__file__)
NUM_TEST = 1000


""" from pt to onnx """

# get pt model
path_pt = os.path.join(PATH, "cnn.pt")
model_pt = torch.load(f=path_pt)

# dummy input
dummy_input = torch.randn(size=(1, 1, 28, 28))

# save onnx model
path_onnx = os.path.join(PATH, "cnn.onnx")
input_names = ['actual_input'] + ['learned_%d' % i for i in range(6)]
torch.onnx.export(model=model_pt, args=dummy_input, f=path_onnx,
                  verbose=True, input_names=input_names)  # arg verbose: True to print log


""" load onnx model """

# load onnx model
model_onnx = onnx.load(path_onnx)

# check if model well formed
onnx.checker.check_model(model_onnx)

# print a human readable representation of the graph
print(onnx.helper.printable_graph(model_onnx.graph))

# data input
test_data = torchvision.datasets.MNIST(root='./data', train=False)
test_data_x = torch.unsqueeze(input=test_data.test_data, dim=1).type(torch.FloatTensor)[: NUM_TEST] / 255.
test_data_y = test_data.test_labels[: NUM_TEST]


""" run onnx model """

# ort session initialize
ort_session = ort.InferenceSession(path_onnx)

# dummy input
outputs = ort_session.run(output_names=None,
                          input_feed={
    
    'actual_input': np.random.randn(1, 1, 28, 28).astype(np.float32)})
print("result of dummy input:  {}".format(outputs[0]), '\n')

# test data, loop
num_correct = 0
for i in range(NUM_TEST):
    test_data_x_, test_data_y_ = test_data_x[i: i + 1], test_data_y[i]
    outputs = ort_session.run(output_names=None, input_feed={
    
    'actual_input': test_data_x_.numpy()})
    predict_y = np.argmax(outputs[0])
    if predict_y == test_data_y_:
        num_correct += 1
    else:
        print("predict result {}, correct answer {}".format(predict_y, test_data_y_), '\n')
accuracy = round(num_correct / NUM_TEST, 3)
print("model accuracy:  {}".format(accuracy), '\n')

Vorsichtsmaßnahmen:

1. Die Modellkonvertierung von ONNX erfordert die Modellstruktur des ursprünglichen Modells;

Zweitens wird der im vorherigen Code gezeigte Prozess zum Ausführen des Modells durch wiederholtes Aufrufen des Modells für alle Testdaten auf zyklische Weise realisiert.Tatsächlich kann dies durch Batch-Input realisiert werden.
Legen Sie zunächst die erste Dimension der Größe der Dummy-Eingabe fest, die zum Erstellen des ONNX-Modells auf die Stapelgröße verwendet wird:

dummy_input = torch.randn(size=(NUM_TEST, 1, 28, 28))  # batch mode

Stellen Sie dann die Länge der ONNX-Modelleingangsdaten entsprechend ein:

# test data, batch mode
test_data_x_, test_data_y_ = test_data_x[: NUM_TEST], test_data_y[: NUM_TEST]
outputs = ort_session.run(output_names=None, input_feed={
    
    'actual_input': test_data_x_.numpy()})
output = outputs[0]
num_correct = 0
for i in range(len(output)):
    predict_y = np.argmax(output[i])
    if predict_y == test_data_y_[i]:
        num_correct += 1
    else:
        print("predict result {}, correct answer {}".format(predict_y, test_data_y_[i]), '\n')
accuracy = round(num_correct / NUM_TEST, 3)
print("model accuracy:  {}".format(accuracy), '\n')

Der Fall, in dem die Modelleingabe aus mehreren Parametern besteht

Im tatsächlichen Modell des maschinellen Lernens kann es mehr als einen Eingabeparameter geben. Um die Multi-Input-Situation zu verifizieren, kopiert dieser Artikel die Modellstruktur des ursprünglichen CNN-Modells, um zwei parallele Netzwerke zu bilden, und mittelt schließlich die Ausgabe der zwei Netzwerke Wert.

Die Modellstruktur ist wie folgt:

import torch.nn as nn


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()

        self.conv11 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # for 28*28 to 14*14
        )
        self.conv12 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # for 14*14 to 7*7
        )
        self.output1 = nn.Linear(32 * 7 * 7, 10)

        self.conv21 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # for 28*28 to 14*14
        )
        self.conv22 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(3, 3), stride=(1, 1), padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # for 14*14 to 7*7
        )
        self.output2 = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x1, x2):
        out1 = self.conv11(x1)
        out1 = self.conv12(out1)
        out1 = out1.view(out1.size(0), -1)  # flatten the 32*7*7 dimension to 1568
        out1 = self.output1(out1)

        out2 = self.conv21(x2)
        out2 = self.conv22(out2)
        out2 = out2.view(out2.size(0), -1)  # flatten the 32*7*7 dimension to 1568
        out2 = self.output1(out2)

        out = (out1 + out2) / 2

        return out

Der Code für das Trainieren und Testen von Modellen lautet wie folgt:

import torch
import torch.nn as nn
import torch.utils.data as data
import torchvision
import matplotlib.pyplot as plt

from model_2 import CNN


# global variables
EPOCH = 5
BATCH_SIZE = 50
LR = 0.001  # learning rate
DOWNLOAD_MNIST = False  # if need to download data, set True at first time


""" data process """

# read train data
train_data = torchvision.datasets.MNIST(
    root='./data', train=True, download=DOWNLOAD_MNIST, transform=torchvision.transforms.ToTensor())
print()
print("size of train_data.train_data:  {}".format(train_data.train_data.size()))  # train_data.train_data is a Tensor
print("size of train_data.train_labels:  {}".format(train_data.train_labels.size()), '\n')
plt.imshow(train_data.train_data[50].numpy(), cmap='Greys')  # plot one example
plt.title('{}'.format(train_data.train_labels[50]))
# plt.show()

# data loader
# combines a dataset and a sampler, and provides an iterable over the given dataset
train_loader = data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)

# read test data
test_data = torchvision.datasets.MNIST(root='./data', train=False)
num_test, num_test_over = 2000, 1000
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(
    torch.FloatTensor)[: num_test] / 255.  # unsqueeze because of 1 channel; value in range (0, 1)
test_y = test_data.test_labels[: num_test]
test_over_x = torch.unsqueeze(
    test_data.test_data, dim=1).type(torch.FloatTensor)[-num_test_over:] / 255.  # test data after training
test_over_y = test_data.test_labels[-num_test_over:]


""" train """

cnn = CNN()
print("CNN model structure:\n")
print(cnn, '\n')
optimizer = torch.optim.Adam(params=cnn.parameters(), lr=LR)  # optimizer: Adam
loss_func = nn.CrossEntropyLoss()  # loss function: cross entropy

for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):  # x, y: Tensor of input and output
        output = cnn(x, x)
        loss = loss_func(output, y)

        # optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # test at each 50 steps
        if not step % 50:
            test_output = cnn(test_x, test_x)
            predict_y = torch.max(test_output, 1)[1].data.numpy()
            accuracy = float((predict_y == test_y.data.numpy()).astype(int).sum()) / float(test_y.size(0))
            print('Epoch: ', epoch, '| train loss: %.4f' % loss.data.numpy(), '| test accuracy: %.3f' % accuracy)


# output model file
torch.save(cnn, 'cnn2.pt')
print()
print('finish training', '\n')


""" test """

# load model
print('load cnn model', '\n')
cnn_ = torch.load('cnn2.pt')

# test new data
test_output = cnn_(test_over_x, test_over_x)
predict_y = torch.max(test_output, 1)[1].data.numpy()
print("prediction number:  {}".format(predict_y))
print("real number:  {}".format(test_over_y.numpy()), '\n')
accuracy = float((predict_y == test_over_y.data.numpy()).astype(int).sum()) / float(test_over_y.size(0))
print("accuracy:  {}".format(accuracy), '\n')

ONNX-Modellkonvertierung und -betrieb:
(Hier ist der Batch-Betriebsmodus, achten Sie auf die Eingabenamenseinstellung mehrerer Eingaben)

import os

import numpy as np
import torch
import torchvision
import onnx
import onnxruntime as ort

from model_2 import CNN  # need model structure


PATH = os.path.dirname(__file__)
NUM_TEST = 1000


""" from pt to onnx """

# get pt model
path_pt = os.path.join(PATH, "cnn2.pt")
model_pt = torch.load(f=path_pt)

# dummy input
dummy_input = torch.randn(size=(NUM_TEST, 1, 28, 28))  # batch mode

# save onnx model
path_onnx = os.path.join(PATH, "cnn2.onnx")
input_names = ['actual_input_{}'.format(i) for i in range(2)] + ['learned_%d' % i for i in range(8)]  # 2 inputs
torch.onnx.export(model=model_pt, args=(dummy_input, dummy_input), f=path_onnx,
                  verbose=True, input_names=input_names)  # arg verbose: True to print log


""" check onnx model """

# load onnx model
model_onnx = onnx.load(path_onnx)

# check if model well formed
onnx.checker.check_model(model_onnx)

# print a human readable representation of the graph
print(onnx.helper.printable_graph(model_onnx.graph))

# data input
test_data = torchvision.datasets.MNIST(root='./data', train=False)
test_data_x = torch.unsqueeze(input=test_data.test_data, dim=1).type(torch.FloatTensor)[: NUM_TEST] / 255.
test_data_y = test_data.test_labels[: NUM_TEST]


""" run onnx model """

# ort session initialize
ort_session = ort.InferenceSession(path_onnx)

# test data, batch mode
test_data_x_, test_data_y_ = test_data_x[: NUM_TEST], test_data_y[: NUM_TEST]
outputs = ort_session.run(output_names=None, input_feed={
    
    'actual_input_0': test_data_x_.numpy(),
                                                         'actual_input_1': test_data_x_.numpy()})
output = outputs[0]
num_correct = 0
for i in range(len(output)):
    predict_y = np.argmax(output[i])
    if predict_y == test_data_y_[i]:
        num_correct += 1
    else:
        print("predict result {}, correct answer {}".format(predict_y, test_data_y_[i]), '\n')
accuracy = round(num_correct / NUM_TEST, 3)
print("model accuracy:  {}".format(accuracy), '\n')

Vergleich des PyTorch-Modells und des ONNX-Modells

Wie bereits erwähnt, zeigt die Erfahrung, dass die Betriebseffizienz des ONNX-Modells deutlich besser ist als die des ursprünglichen PyTorch-Modells. Dies scheint auf die Optimierung im ONNX-Modellgenerierungsprozess zurückzuführen zu sein, die auch den Modellgenerierungsprozess zeitaufwändiger macht. aufwendig, aber die Gesamteffizienz ist immer noch beachtlich. .

Darüber hinaus ist gemäß der statistischen Analyse der Laufergebnisse des ONNX-Modells und des PyTorch-Modells (Mittelwert und Standardabweichung des Fehlers) ersichtlich, dass die Laufergebnisse des ONNX-Modells sehr kleine Fehler aufweisen und im Wesentlichen sind zuverlässig.