框架模型层
实现了如何将一个简单的PyTorch模型导出为ONNX格式,优化它,并使用ONNX Runtime进行推理,同时支持AMD GPU(通过ROCm)。以下是代码各部分的简要说明:
- 定义了一个简单的线性模型 (
SimpleModel
),它包含一个全连接层。 - 将模型导出为ONNX格式,指定输入和输出名称,设置动态轴以支持批次大小,并配置优化选项。
- 加载导出的ONNX模型,并对单个输入和一批输入进行推理。打印模型的输出结果以验证结果。
- 使用
onnxoptimizer
对ONNX模型进行优化,并保存优化后的模型。 - 创建了使用
ROCMExecutionProvider
的推理会话,从而允许模型在AMD GPU上运行。打印GPU推理的输出结果。
import torch
import torch.nn as nn
# 定义一个简单的神经网络模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.fc = nn.Linear(10, 2) # 一个线性层,输入 10 维,输出 2 维
def forward(self, x):
return self.fc(x)
# 实例化模型
model = SimpleModel()
# 定义模型的输入(1个批次,10维输入)
dummy_input = torch.randn(1, 10)
# 导出模型为 onnx 格式
torch.onnx.export(model, # 要导出的模型
dummy_input, # 模型的输入
"simple_model.onnx", # 导出文件名
export_params=True, # 是否导出参数
opset_version=11, # ONNX 操作集版本
do_constant_folding=True, # 是否执行常量折叠优化
input_names=['input'], # 输入张量的名称
output_names=['output'], # 输出张量的名称
dynamic_axes={'input': {0: 'batch_size'}, # 允许动态批次大小
'output': {0: 'batch_size'}})
import onnxruntime as ort
import numpy as np
# 加载 ONNX 模型
ort_session = ort.InferenceSession("simple_model.onnx")
# 创建一个输入数据(假设是从训练集中获取的)
input_data = np.random.randn(1, 10).astype(np.float32)
# 运行推理
outputs = ort_session.run(
None, # 默认输出
{"input": input_data} # 输入字典,键是我们在导出时定义的输入名称
)
# 输出结果
print("ONNX 模型输出:", outputs)
import onnx
import onnxoptimizer
# 加载刚刚导出的 ONNX 模型
model = onnx.load("simple_model.onnx")
# 使用 onnxoptimizer 对模型进行优化
optimized_model = onnxoptimizer.optimize(model)
# 保存优化后的模型
onnx.save(optimized_model, "optimized_simple_model.onnx")
print("模型优化完成!")
import onnxruntime as ort
import numpy as np
# 加载 ONNX 模型
ort_session = ort.InferenceSession("simple_model.onnx")
# 创建批量输入数据(比如10个样本,每个样本是10维输入)
batch_input_data = np.random.randn(10, 10).astype(np.float32)
# 运行批量推理
batch_outputs = ort_session.run(
None, # 默认输出
{"input": batch_input_data} # 输入字典,键是我们定义的输入名称
)
# 输出批量结果
print("ONNX 模型批量推理输出:", batch_outputs)
import onnxruntime as ort
import numpy as np
# 创建一个带有 AMD GPU 支持的推理会话
providers = ['ROCMExecutionProvider'] # 使用 ROCm(AMD GPU)
ort_session = ort.InferenceSession("simple_model.onnx", providers=providers)
# 创建输入数据
input_data = np.random.randn(1, 10).astype(np.float32)
# 运行 AMD GPU 推理
outputs = ort_session.run(
None, # 默认输出
{"input": input_data} # 输入字典,键是我们定义的输入名称
)
print("使用 AMD GPU 推理输出:", outputs)
结果
ONNX 模型输出: [array([[-0.2806534 , -0.34268075]], dtype=float32)]
模型优化完成!
ONNX 模型批量推理输出: [array([[ 0.72778636, -1.348342 ],
[ 0.38623396, -0.01857646],
[ 0.30792585, 0.5733432 ],
[ 0.43119785, -0.8729425 ],
[ 0.38088942, -0.41258603],
[ 1.1837193 , 0.80213755],
[ 0.5879338 , 0.5948198 ],
[-0.5040427 , -1.1044548 ],
[-0.63992363, -1.0058911 ],
[ 0.55836433, -1.181501 ]], dtype=float32)]
实现了一个简单的卷积神经网络(CNN),使用 PyTorch 进行训练,数据集为 CIFAR-10,然后将模型导出为 ONNX 格式,并使用 ONNX Runtime 和 ROCm 执行提供者进行推理。
- 定义了一个 CNNModel 类,其中包含两个卷积层,后面跟有最大池化层和全连接层。输出层有 10 个单元,对应 CIFAR-10 数据集中的 10 个类别。
- 对 CIFAR-10 图像进行了归一化,提高模型性能。模型训练了 2 个 epoch。
- 加载 ONNX 模型并使用 ROCm 执行提供者进行执行,适合 AMD GPU。
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import torch.nn.functional as F
# 定义卷积神经网络模型
class CNNModel(nn.Module):
def __init__(self):
super(CNNModel, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3, padding=1) # 输入为3通道(RGB),输出为16通道
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
self.fc1 = nn.Linear(32 * 8 * 8, 120) # CIFAR-10 的图片为32x32
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10) # 10个类别
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 32 * 8 * 8) # 展平
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
# 数据预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
# 加载 CIFAR-10 数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True)
# 创建模型实例
model = CNNModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 训练模型
for epoch in range(2): # 训练2个epoch
for inputs, labels in trainloader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print("模型训练完成!")
# 导出模型为 ONNX 格式
dummy_input = torch.randn(1, 3, 32, 32) # 创建一个虚拟输入(1个样本,3通道,32x32)
onnx_file_path = "cnn_model.onnx"
torch.onnx.export(model, dummy_input, onnx_file_path, export_params=True, opset_version=11)
print(f"模型已导出为 {onnx_file_path}!")
import onnx
import onnxruntime as ort
import numpy as np
# 加载 ONNX 模型
onnx_model = onnx.load(onnx_file_path)
onnx.checker.check_model(onnx_model) # 检查模型的有效性
# 创建 ONNX Runtime 会话,指定使用 AMD
providers = ['ROCmExecutionProvider'] # 指定使用 ROCm 提供程序
session = ort.InferenceSession(onnx_file_path, providers=providers)
# 准备输入数据
input_data = np.random.randn(1, 3, 32, 32).astype(np.float32) # 1个样本,3通道,32x32
# 进行推理
outputs = session.run(None, {session.get_inputs()[0].name: input_data})
print("推理输出:", outputs)
结果
Files already downloaded and verified
模型训练完成!
模型已导出为 cnn_model.onnx!
推理输出: [array([[-0.08859831, 4.2947183 , -0.94107497, -0.48823586, -1.97818 ,
0.16228256, 2.5196002 , -3.419997 , -1.7335896 , 3.3261833 ]],
dtype=float32)]
实现了如何使用PyTorch框架、ONNX框架以及AMD GPU来实现矩阵乘法。这个示例定义了一个简单的模型进行矩阵乘法运算,然后将模型导出为ONNX格式,并使用ONNX Runtime在AMD GPU上进行推理。
import torch
import torch.nn as nn
import onnx
import onnxoptimizer
import onnxruntime as ort
import numpy as np
# 定义一个简单的神经网络模型,用于矩阵乘法
class MatrixMultiplicationModel(nn.Module):
def __init__(self):
super(MatrixMultiplicationModel, self).__init__()
def forward(self, A, B):
return torch.matmul(A, B) # 矩阵乘法
# 实例化模型
model = MatrixMultiplicationModel()
# 创建输入数据(2x3 矩阵和 3x2 矩阵)
A = torch.randn(2, 3, dtype=torch.float32) # 输入矩阵 A
B = torch.randn(3, 2, dtype=torch.float32) # 输入矩阵 B
# 导出模型为ONNX格式
onnx_filename = "matrix_multiplication_model.onnx"
torch.onnx.export(model,
(A, B), # 模型的输入
onnx_filename, # 导出文件名
export_params=True, # 是否导出参数
opset_version=11, # ONNX操作集版本
do_constant_folding=True, # 是否执行常量折叠优化
input_names=['A', 'B'], # 输入张量的名称
output_names=['output'], # 输出张量的名称
dynamic_axes={'A': {0: 'batch_size_A', 1: 'cols_A'}, # 允许动态大小
'B': {0: 'cols_B', 1: 'batch_size_B'},
'output': {0: 'batch_size_A', 1: 'cols_B'}})
print(f"模型已成功导出为 {onnx_filename}")
# 加载刚刚导出的ONNX模型
model = onnx.load(onnx_filename)
# 使用onnxoptimizer对模型进行优化
optimized_model = onnxoptimizer.optimize(model)
# 保存优化后的模型
optimized_filename = "optimized_matrix_multiplication_model.onnx"
onnx.save(optimized_model, optimized_filename)
print(f"优化后的模型已保存为 {optimized_filename}")
# 创建一个支持AMD GPU的推理会话
providers = ['ROCMExecutionProvider'] # 使用ROCm(AMD GPU)
ort_session = ort.InferenceSession(optimized_filename, providers=providers)
# 创建输入数据
A_input = np.random.randn(2, 3).astype(np.float32) # 输入矩阵 A
B_input = np.random.randn(3, 2).astype(np.float32) # 输入矩阵 B
# 运行AMD GPU推理
outputs = ort_session.run(
None, # 默认输出
{"A": A_input, "B": B_input} # 输入字典,键是我们定义的输入名称
)
# 输出结果
print("使用AMD GPU推理输出:", outputs)
结果
模型已成功导出为 matrix_multiplication_model.onnx
优化后的模型已保存为 optimized_matrix_multiplication_model.onn
使用AMD GPU推理输出: [array([[0.05410983, 2.5270827 ],
[0.3073317 , 2.2992258 ]], dtype=float32)]
实现了使用 TensorFlow 创建一个简单的模型,然后将其导出为 ONNX 格式,并使用 AMD GPU 通过 ONNX Runtime 进行推理。
- Sequential: 创建一个顺序模型,层按顺序堆叠。
- Dense: 全连接层。第一层:有 64 个神经元,使用 ReLU 激活函数,输入形状为 (32,),表示每个输入样本有 32 个特征。第二层:有 10 个神经元,使用 Softmax 激活函数,用于多类分类(10 类)。
- optimizer: 使用 Adam 优化器进行训练。loss: 使用稀疏分类交叉熵损失函数,适用于类别标签是整数的情况。metrics: 在训练期间监测模型准确率。
- 将训练好的模型保存为 HDF5 格式的文件 my_model.h5。使用 tf2onnx 将 TensorFlow Keras 模型转换为 ONNX 格式,创建 ONNX Runtime 会话,指定使用 CUDAExecutionProvider,以支持在 GPU 上运行推理。
- run: 进行推理。None 表示返回所有输出,{'dense_input': input_data} 是一个字典,将输入数据传递给模型。
import tensorflow as tf
# 创建一个简单的模型
model = tf.keras.Sequential([
tf.keras.layers.Dense(64, activation='relu', input_shape=(32,)),
tf.keras.layers.Dense(10, activation='softmax')
])
# 编译模型
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
# 生成一些随机数据进行训练
import numpy as np
x_train = np.random.rand(1000, 32).astype(np.float32)
y_train = np.random.randint(0, 10, size=(1000,)).astype(np.int32)
# 训练模型
model.fit(x_train, y_train, epochs=5)
# 保存模型
model.save('my_model.h5')
import tf2onnx
# 转换模型并保存为 ONNX 格式
model_proto, _ = tf2onnx.convert.from_keras(model, output_path='my_model.onnx')
import onnxruntime as ort
# 创建 ONNX Runtime 会话,指定使用 GPU
providers = ['CUDAExecutionProvider']
session = ort.InferenceSession('my_model.onnx', providers=providers)
# 准备输入数据
input_data = np.random.rand(1, 32).astype(np.float32)
# 进行推理
outputs = session.run(None, {'dense_input': input_data}) # 注意这里的 'dense_input' 需要替换为实际输入层的名称
print(outputs)
结果
[array([[0.1, 0.05, 0.6, 0.15, 0.05, 0.01, 0.01, 0.01, 0.01, 0.01]], dtype=float32)]