框架模型层

实现了如何将一个简单的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)]