MNISTをCUDAで学習して、DirectMLで推論する

古典的な、機械学習のサンプルであるMNISTをCUDAで学習して、DirectMLで推論してみます。

まず、CUDAで学習するコードを示します。

# %%
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms

# ハイパーパラメータなどの設定値
num_epochs = 10  # 学習を繰り返す回数
num_batch = 100  # 一度に処理する画像の枚数
learning_rate = 0.001  # 学習率
image_size = 28 * 28  # 画像の画素数(幅x高さ)

# GPU(CUDA)が使えるかどうか？
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 学習用／評価用のデータセットの作成
transform = transforms.Compose([transforms.ToTensor()])

# MNISTデータの取得
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, transform=transform)

# データローダー
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=num_batch, shuffle=True)
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=num_batch, shuffle=True)

# ニューラルネットワークモデルの定義
class Net(nn.Module):
    def __init__(self, input_size, output_size):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(input_size, 100)
        self.fc2 = nn.Linear(100, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.sigmoid(x)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

# ニューラルネットワークの生成
model = Net(image_size, 10).to(device)

# 損失関数の設定
criterion = nn.CrossEntropyLoss()

# 最適化手法の設定
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 学習
model.train()
for epoch in range(num_epochs):
    loss_sum = 0
    for inputs, labels in train_dataloader:
        inputs = inputs.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        inputs = inputs.view(-1, image_size)
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        loss_sum += loss.item()
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss_sum / len(train_dataloader):.4f}")

# ONNX形式でモデルをエクスポート
dummy_input = torch.randn(1, image_size).to(device)
torch.onnx.export(model, dummy_input, "../Models/mnist_model.onnx", verbose=True)

print("モデルをONNX形式でエクスポートしました。")

このコードを実行すると、ONNX形式のモデルファイルが得られます。ここで得た、モデルファイルを用いてDirectMLで推論します。今回は、ONNX Runtimeを用いてみます。

# %%
import onnxruntime as ort
import numpy as np

# ONNXモデルの読み込み
onnx_model_path = "../Models/mnist_model.onnx"
session = ort.InferenceSession(onnx_model_path)

# 入力データの生成（例：手書き数字の画像）
input_data = np.random.rand(1, 784).astype(np.float32)  # 28x28の画像を1次元に変換

# 推論の実行
output = session.run(None, {"onnx::Gemm_0": input_data})

# 結果の表示
predicted_class = np.argmax(output[0])
print(f"Predicted class: {predicted_class}")

推論コードではランダムに28×28ピクセルの画像を生成して、モデルで推論を行っています。

当初はpytorch-directmlでMNISTの学習側も出来ないかなと考慮しましたが、テンソルの演算などは実装可能なようですが、MNISTの学習側のコードを作ると、Runtime Errorを引き起こし、動かないので、pytorch-directmlのissuesを見たところ、学習側を作るには足りていない機能があるようなので、一旦、そちらの方向はペンディングにし、学習側は一旦、通常のpytorchの作りにしました。

差し当たり、想定した、CUDAでモデルを作って、DirectMLで推論するというタスクは可能だと確認できました。