PyTorch で複数モデルを繋げる方法

PyTorchで複数のモデルの出力をアンサンブルする方法

こんにちは。今日は、PyTorchで複数の深層学習モデルの出力をアンサンブルする方法について説明します。 アンサンブル学習は、単一のモデルよりも高い予測精度を得るための強力な手法です。異なる種類のモデルや同じモデルの異なる初期化からの複数の出力を組み合わせることで、モデルの一般化性能を向上させることができます。 今回は基礎編ということで非常にシンプルな例を解説します。

ここでは、2つの畳み込みニューラルネットワーク(CNN)の出力を結合し、全結合層に入力して損失を計算する例を示します。

※当然同じアーキテクチャでは多様性が獲得されませんが、今回は実装面の簡便さを重視します

ステップ1: ライブラリをインポートする

import torch
import torch.nn as nn
import torch.nn.functional as F

ステップ2: CNNモデルを定義する

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32 * 7 * 7, 64)
        self.fc2 = nn.Linear(64, 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 * 7 * 7)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

この例では、2つの畳み込み層と2つの全結合層からなるCNNモデルを定義しています。入力は3チャンネルの画像で、出力は10クラスの確率分布です。

ステップ3: 2つのCNNモデルを作成する

model1 = CNN()
model2 = CNN()

ステップ4: 2つのモデルの出力を結合する

def ensemble_models(x):
    output1 = model1(x)
    output2 = model2(x)
    output = torch.cat((output1, output2), dim=1)
    return output

この関数では、2つのCNNモデルの出力を結合しています。torch.catを使用して、2つのテンソルを次元1(列方向)に沿って結合しています。つまり、2つの出力を単に連結しています。

ステップ5: 全結合層を定義する

fc = nn.Linear(20, 10)

ここでは、20次元の入力(2つのCNNモデルの出力の結合)を受け取り、10クラスに分類する全結合層を定義しています。

ステップ6: 損失関数と最適化アルゴリズムを定義する

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(list(model1.parameters()) + list(model2.parameters()) + list(fc.parameters()), lr=0.001)

交差エントロピー損失関数とSGD最適化アルゴリズムを定義しています。最適化アルゴリズムには、2つのCNNモデルと全結合層のパラメータをすべて含めています。

ステップ7: 学習とテストのループ

for epoch in range(num_epochs):
    for data, labels in train_loader:
        optimizer.zero_grad()
        outputs = fc(ensemble_models(data))
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    # テストデータでモデルを評価
    correct = 0
    total = 0
    with torch.no_grad():
        for data, labels in test_loader:
            outputs = fc(ensemble_models(data))
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'Epoch [{epoch+1}/{num_epochs}], Accuracy: {accuracy:.2f}%')

学習ループでは、2つのCNNモデルの出力をensemble_models関数で結合し、全結合層に入力して損失を計算しています。その後、逆伝播と最適化を行います。 テストループでは、テストデータに対する予測を行い、正解率を計算しています。