関連ドキュメントは PyTorch 公式 Web サイトで提供されています。興味のある学生は次のドキュメントを参照してください: EXPORTING A MODEL FROM PYTORCH TO ONNX AND RUNNING IT USING ONNX RUNTIME
1. 準備
- セマンティック セグメンテーション モデル
model.py - 訓練された重みファイル
model.pth / model.pt onnx==1.12.0onnxruntime==1.15.1
import torch.onnx
from models import PPLiteSeg
import onnxruntime as ort
from PIL import Image
import numpy as np
import torchvision.transforms as transforms
2.PyTorchモデルの作成
まず、PyTorch モデルを作成し、.pth重みファイルをロードする必要があります。
# 创建模型
torch_model = PPLiteSeg()
# 加载模型权重
model_state_dict = torch.load("checkpoint/model.pth")
# 如果模型使用了DDP训练，则模型状态字典的会有'module'的前缀，我们需要删除
# 创建一个新的字典，去掉 "module." 前缀
# new_state_dict = {k.replace('module.', ''): v for k, v in model_state_dict['model'].items()}
# 加载模型权重
torch_model.load_state_dict(new_state_dict, strict=True)
print("\033[1;31m模型权重加载完毕...\033[0m")
"""
因为我们的模型最终的输出并没有经过后处理，此时的shape为[N, num_classes, H, W]，所以需要对模型添加上后处理,
让模型的输出为[N, 1, H, W]
"""
# 给模型添加后处理操作
torch_model = WrappedModel(torch_model)
# 设置模型为推理状态(这一步是必须的!)
torch_model.eval()
# 创建一个输入Tensor
x = torch.randn(1, 3, 512, 512, requires_grad=True)
torch_out = torch_model(x)
print(torch_out[0].shape) # torch.Size([1, 1, 512, 512])
コードWrappedModelは次のとおりです。
import torch
class WrappedModel(torch.nn.Module):
def __init__(self, model, output_op):
super().__init__()
self.model = model
def forward(self, x):
outs = self.model(x)
new_outs = []
for out in outs:
out = torch.nn.functional.softmax(out, dim=1) # 沿着通道维度进行概率计算
label = torch.argmax(out, dim=1).to(dtype=torch.int32) # 获取最大的位置
label = torch.unsqueeze(label, 1)
# torch.max返回值有两个：最大值的张量 + 最大值的索引张量
max_score = torch.max(out, dim=1)[0] # 获取最大概率
max_score = torch.unsqueeze(max_score, 1)
new_outs.append(label)
new_outs.append(max_score)
# 返回的是一个len==2的list
return new_outs
この時点で、PyTorch モデルが正常に作成され、トレーニングされた重みが正しく読み込まれたことを意味します。
3. ONNX モデルに変換して保存します
# Export the model
torch.onnx.export(torch_model, # model being run
x, # model input (or a tuple for multiple inputs)
"model.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=11, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['input'], # the model's input names
output_names = ['label', 'score'], # the model's output names
dynamic_axes={
'input' : {
0 : 'B'}, # variable length axes
'output' : {
0 : 'B'}})
print("\033[1;31mONNX模型转换完毕.\033[0m")
以下は、torch.onnx.export関数のパラメータの説明です。
-
torch_model: これは、エクスポートされる PyTorch モデルのインスタンスです。 -
x: これはモデルの入力データであり、モデルの入力方法に応じて、単一の入力 Tensor または複数の入力 Tensor を含むタプルにすることができます。 -
"model.onnx": エクスポートされた ONNX モデル ファイルの保存パスです。ONNX モデルは「model.onnx」という名前のファイルに保存されます。ファイル名とパスは変更できます。 -
export_params=True: これは、モデルのパラメータの重みをエクスポートするかどうかを示すブール値です。に設定するとTrue、モデルのパラメータはモデルとともに ONNX ファイルに保存され、推論時に使用されます。に設定するとFalse、パラメータはエクスポートされず、モデル構造のみがエクスポートされます。 -
opset_version=11: これは、モデルのエクスポートに使用される ONNX バージョンです。この例では、ONNX バージョン 11 が使用されます。ONNX のバージョンが異なればサポートされる操作も異なるため、モデルとランタイムと互換性のあるバージョンを選択する必要があります。 -
do_constant_folding=True: これは、最適化のために定数折りたたみを実行するかどうかを示すブール値です。に設定するとTrue、ONNX エクスポートは、モデル内の定数テンソルを定数ノードに折りたたんで、モデル ファイル サイズを削減し、推論速度を向上させようとします。 -
input_nameslist: これは、モデルの入力 Tensor を識別する入力名 ( ) のモデルのリストです。この例では、モデルの入力 Tensor の名前は「input」です。 -
output_nameslist: これは、モデルの出力 Tensor を識別するモデルの出力名 ( ) のリストです。この例では、モデルの出力テンソルの名前は「label」と「score」です。 -
dynamic_axes: 動的軸の名前を指定する辞書です。動的軸は可変長を持つことができる軸であり、通常はバッチ軸です。この例では、入力「input」と出力「output」の最初の次元が「B」として指定されており、バッチ軸が可変長であることを示しています。
これらのパラメーターを使用すると、PyTorch モデルを ONNX 形式にエクスポートする方法を制御し、ニーズに応じて構成できます。
説明:
- このモデルの出力は長さ 2 のリストであるため、
output_names2 つあるはずです。 dynamic_axesここでは、バッチ ディメンションを動的に設定します。つまり、ONNX モデルのバッチ ディメンションの入力は任意であり、固定されていません。
完全なコードは次のとおりです。
import torch
import numpy as np
import torch.onnx
from models import PPLiteSeg
class WrappedModel(torch.nn.Module):
def __init__(self, model, output_op):
super().__init__()
self.model = model
def forward(self, x):
outs = self.model(x)
new_outs = []
for out in outs:
out = torch.nn.functional.softmax(out, dim=1) # 沿着通道维度进行概率计算
label = torch.argmax(out, dim=1).to(dtype=torch.int32) # 获取最大的位置
label = torch.unsqueeze(label, 1)
# torch.max返回值有两个：最大值的张量 + 最大值的索引张量
max_score = torch.max(out, dim=1)[0] # 获取最大概率
max_score = torch.unsqueeze(max_score, 1)
new_outs.append(label)
new_outs.append(max_score)
# 返回的是一个len==2的list
return new_outs
if __name__ == "__main__":
# 创建模型
torch_model = PPLiteSeg()
# 加载模型权重
model_state_dict = torch.load("checkpoint/model.pth")
# 如果模型使用了DDP训练，则模型状态字典的会有'module'的前缀，我们需要删除
# 创建一个新的字典，去掉 "module." 前缀
# new_state_dict = {k.replace('module.', ''): v for k, v in model_state_dict['model'].items()}
# 加载模型权重
torch_model.load_state_dict(new_state_dict, strict=True)
print("\033[1;31m模型权重加载完毕...\033[0m")
"""
因为我们的模型最终的输出并没有经过后处理，此时的shape为[N, num_classes, H, W]，所以需要对模型添加上后处理,
让模型的输出为[N, 1, H, W]
"""
# 给模型添加后处理操作
torch_model = WrappedModel(torch_model)
# 设置模型为推理状态(这一步是必须的!)
torch_model.eval()
# 创建一个输入Tensor
x = torch.randn(1, 3, 512, 512, requires_grad=True)
torch_out = torch_model(x)
# Export the model
torch.onnx.export(torch_model, # model being run
x, # model input (or a tuple for multiple inputs)
"model.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=11, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['input'], # the model's input names
output_names = ['label', 'score'], # the model's output names
dynamic_axes={
'input' : {
0 : 'B'}, # variable length axes
'output' : {
0 : 'B'}})
print("\033[1;31mONNX模型转换完毕.\033[0m")
4.ONNXを変更する
4.1 入力と出力の形状を変更する
ONNX として保存した後、以下に示すように、 Netronというソフトウェアを使用して.onnxファイルを開くことができます。
ONNX ファイル内のArgMax対応する出力がであることがわかりlabel、モデル変換が正しいことを示しています。しかし、右側を見ると、形状は であることがわかります。これは私たちが望むものですが、出力は論理的には同じであるはずですが、そうではありません。後の TRT (TensorRT) への変換を容易にするために、出力を変更します。変更後のコードは次のとおりです。ReduceMaxscoreinput[B, 3, 512, 512][B, 3, 512, 512]
import onnx
import argparse
def show_inp_and_oup_info(model, modify=False):
input_info = model.graph.input
print("模型的输入信息：")
for info in input_info:
print(info.name, info.type)
output_info = model.graph.output
print("模型的输出信息：")
for info in output_info:
print(info.name, info.type)
if __name__ == "__main__":
# 输入 ONNX 模型路径
model_path = "model.onnx"
# 输出 ONNX 模型路径
output_path = "retype_model.onnx"
# 读取 ONNX 模型
model = onnx.load(model_path)
show_inp_and_oup_info(model, modify=False)
# 找到输入张量并修改
# for input_info in model.graph.input:
# if input_info.name in ['x', 'input']:
# # 修改输入张量的形状
# input_info.type.tensor_type.shape.dim[0].dim_param = "B"
# 修改输出张量的形状
for output_info in model.graph.output:
if output_info.name in ["label", "score"]:
output_info.type.tensor_type.shape.dim[0].dim_param = "B"
output_info.type.tensor_type.shape.dim[2].dim_value = 512
output_info.type.tensor_type.shape.dim[3].dim_value = 512
show_inp_and_oup_info(model, modify=True)
# 保存修改后的模型
onnx.save(model, output_path)
4.2 名前の変更
入力名と出力名を変更したい場合は、次のスクリプトを使用することもできます。
import argparse
import sys
import onnx
def parse_arguments():
parser = argparse.ArgumentParser()
parser.add_argument('--model', required=True, help='Path of directory saved the input model.')
parser.add_argument('--origin_names', required=True, nargs='+', help='The original name you want to modify.')
parser.add_argument('--new_names', required=True, nargs='+',
help='The new name you want change to, the number of new_names should be same with the number of origin_names')
parser.add_argument('--save_file', required=True, help='Path to save the new onnx model.')
return parser.parse_args()
if __name__ == '__main__':
args = parse_arguments()
model = onnx.load(args.model)
output_tensor_names = set()
for ipt in model.graph.input:
output_tensor_names.add(ipt.name)
for node in model.graph.node:
for out in node.output:
output_tensor_names.add(out)
for origin_name in args.origin_names:
if origin_name not in output_tensor_names:
print("[ERROR] Cannot find tensor name '{}' in onnx model graph.".format(origin_name))
sys.exit(-1)
if len(set(args.origin_names)) < len(args.origin_names):
print("[ERROR] There's dumplicate name in --origin_names, which is not allowed.")
sys.exit(-1)
if len(args.new_names) != len(args.origin_names):
print("[ERROR] Number of --new_names must be same with the number of --origin_names.")
sys.exit(-1)
if len(set(args.new_names)) < len(args.new_names):
print("[ERROR] There's dumplicate name in --new_names, which is not allowed.")
sys.exit(-1)
for new_name in args.new_names:
if new_name in output_tensor_names:
print("[ERROR] The defined new_name '{}' is already exist in the onnx model, which is not allowed.")
sys.exit(-1)
for i, ipt in enumerate(model.graph.input):
if ipt.name in args.origin_names:
idx = args.origin_names.index(ipt.name)
model.graph.input[i].name = args.new_names[idx]
for i, node in enumerate(model.graph.node):
for j, ipt in enumerate(node.input):
if ipt in args.origin_names:
idx = args.origin_names.index(ipt)
model.graph.node[i].input[j] = args.new_names[idx]
for j, out in enumerate(node.output):
if out in args.origin_names:
idx = args.origin_names.index(out)
model.graph.node[i].output[j] = args.new_names[idx]
for i, out in enumerate(model.graph.output):
if out.name in args.origin_names:
idx = args.origin_names.index(out.name)
model.graph.output[i].name = args.new_names[idx]
onnx.checker.check_model(model)
onnx.save(model, args.save_file)
print("[Finished] The new model saved in {}.".format(args.save_file))
print("[DEBUG INFO] The inputs of new model: {}".format([x.name for x in model.graph.input]))
print("[DEBUG INFO] The outputs of new model: {}".format([x.name for x in model.graph.output]))
次のようにコマンドを使用します。
python rename_onnx_model_name.py \
--model model.onnx \
--origin_names x y z \
--new_names x1 y1 z1 \
--save_file new_model.onnx
5. 変換前と変換後の効果をテストする
変換前と変換後の効果をテストするには、次の 2 つのアイデアがあります。
- アイデア 1: 2 つのモデルの出力の違いを比較する - マシン ビュー
- アイデア 2: 2 つのモデルの出力を直接画像に変換 - 肉眼で見る
5.1 2 つのモデルの出力の違いを比較する
PyTorch チュートリアルでは、このメソッドが使用されます。
# compare ONNX Runtime and PyTorch results
np.testing.assert_allclose(torch_res[0].numpy(), onnx_res[0], rtol=1e-03, atol=1e-05)
np.testing.assert_allclose(torch_res[1].detach().numpy(), onnx_res[1], rtol=1e-03, atol=1e-05)
print("\033[1;44mExported model has been tested with ONNXRuntime, and the result looks good!\033[0m")
score私たちのモデルには と があるためlabel、両方をテストする必要があります。
5.2 2 つのモデルの出力を画像に直接変換する
以下では詳細なデモは提供せず、必要な機能のみを提供します。
5.2.1 画像をロードして前処理する
def load_test_img(image_path, target_size=(512, 512)):
# 加载图片
image = Image.open(image_path)
# 调整图片大小为目标大小
image = image.resize(target_size, Image.BILINEAR)
# 使用 torchvision.transforms 将 PIL 图片转换为 PyTorch 张量
transform = transforms.Compose([transforms.ToTensor(), # 转换为张量
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) # 归一化
])
# 应用变换并添加批次维度 [1, C, H, W]
image_tensor = transform(image).unsqueeze(0)
return image_tensor
5.2.2 ONNX モデルのロード
def create_onnx_model(ckpt_path):
import onnxruntime as ort
ort_session = ort.InferenceSession(ckpt_path)
print("\033[1;31mONNX模型创建完毕...\033[0m")
return ort_session
5.2.3 ONNX モデルの実行
onnx_res = onnx_model.run(None, {
"input": [test_img.squeeze(0)]})
ここで何か説明する必要があります:
-
onnx_model.run: これは ONNX モデルを実行する方法です。onnx_modelONNX ランタイムを通じて作成された ONNX モデルのインスタンスです。 -
None: これは、目的の出力名を指定するためのプレースホルダーです。この例では、None出力名を指定しないことを意味するため、ONNX ランタイムはすべての出力を返します。 -
{"input": [test_img.squeeze(0)]}: 入力データの辞書です。ONNX モデルでは通常、入力データを指定するための辞書が必要です。キーは入力名、値は入力データです。ここで、入力名は「input」、対応する入力データは ですtest_img.squeeze(0)。
test_img.squeeze(0): これは、test_imgONNX モデルの入力要件に適合するように、Tensor の最初の次元 (通常はバッチ次元) を圧縮 (削除) します。通常、ONNX モデルの入力テンソルはバッチ ディメンションを想定していないため、.squeeze(0)最初のディメンションを削除して入力データを ONNX モデルと互換性のあるものにします。
このコマンドを実行すると、onnx_resONNX モデルの出力が含まれます。結果は通常、出力 Tensor (1 つであることを覚えておいてくださいlist) のリストであり、各要素はモデル出力に対応します。これらの結果は、モデルの出力に基づいてアクセスし、処理できます。onnx_resこの特定の例では、アプリケーションのシナリオに応じて、使用可能なデータまたはその他の後続の操作に変換するために、さらなる処理が必要になる場合があります。
5.2.4 モデルの結果を画像として保存する
def save_torch_res(torch_res, suffix):
# 转换 PyTorch 张量为 NumPy 数组
torch_res_numpy = torch_res[0].squeeze(0).numpy()
# 如果形状不是 [H, W]，可以进一步调整
print(np.shape(torch_res_numpy))
# 如果形状不是 [H, W]，可以进一步调整
if torch_res_numpy.shape[0] == 1:
torch_res_numpy = torch_res_numpy[0]
# 创建灰度图像
gray_image = Image.fromarray((torch_res_numpy * 255).astype('uint8'), mode='L')
# 将灰度图像转换为伪彩色图像（伪彩色映射可根据需要更改）
pseudo_color_image = gray_image.convert('P', palette=Image.ADAPTIVE, colors=256)
# 保存伪彩色图像
pseudo_color_image.save("results/pytorch_output_pseudo_color_image.png")
print("伪彩色图像已保存为 'results/pytorch_output_pseudo_color_image.png'")
def save_onnx_res(onnx_res, suffix):
# 转换 ONNX 结果为 NumPy 数组
onnx_res_numpy = np.array(onnx_res[0])
# 如果形状不是 [H, W]，可以进一步调整
if onnx_res_numpy.shape[0] == 1:
onnx_res_numpy = np.squeeze(onnx_res_numpy, axis=0)
onnx_res_numpy = np.squeeze(onnx_res_numpy, axis=0)
# 创建灰度图像
gray_image = Image.fromarray((onnx_res_numpy * 255).astype('uint8'), mode='L')
# 将灰度图像转换为伪彩色图像（伪彩色映射可根据需要更改）
pseudo_color_image = gray_image.convert('P', palette=Image.ADAPTIVE, colors=256)
# 保存伪彩色图像
pseudo_color_image.save("results/onnx_output_pseudo_color_image.png")
print("伪彩色图像已保存为 'results/onnx_output_pseudo_color_image.png'")