PyTorchで複数の入力データがある場合の実装方法

IT, Python

この記事では、PyTorchを使用して深層学習モデルに複数の入力を処理する方法について解説します。具体的には、モデルの構築、DataLoaderの作成と使用方法について順を追って説明し、最終的に全体のコードを示します。

スポンサーリンク

モデルの実装方法

ここでは、入力が2か所あるネットワークのサンプルを実装します。最初に、畳み込み層に「高さ15×幅20×チャネル1」の2次元行列を入力し、その出力と「3成分のベクトル」を全結合層に入力する例を示します。

以下はネットワークのコードです：

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.block1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=1),
            nn.BatchNorm2d(32)
        )
        self.block2 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=1),
            nn.BatchNorm2d(64)
        )
        self.full_connection = nn.Sequential(
            nn.Linear(in_features=64*13*18+3, out_features=1024), # '+3' : in_oneDの3ユニット分を追加
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(in_features=1024, out_features=100, bias=False)
        )

    def forward(self, in_twoD, in_oneD):
        x = self.block1(in_twoD)
        x = self.block2(x)
        x = x.view(x.size(0), -1)
        x = torch.cat([x, in_oneD], dim=1)
        y = self.full_connection(x)
        return y

モデルの構築では、__init__メソッドで各層を定義します。特に、全結合層の入力数を畳み込み層からの出力数とin_oneDデータのサイズを合計したものに設定します。

forwardメソッドでは、torch.catメソッドを使って、畳み込み層の出力と追加のベクトル入力を結合します。

DataLoaderの作り方

次に、DataLoaderの作成方法を説明します。複数の入力を扱うための実装に焦点を当て、データの内容はランダムに生成します。

import numpy as np
from sklearn.model_selection import train_test_split
from torch.utils.data import TensorDataset, DataLoader

# サンプルデータの作成
in_oneD, in_twoD, target = [], [], []

for _ in range(100):
    in_oneD.append(np.random.randn(3))
    in_twoD.append(np.random.randn(1, 15, 20))
    target.append(np.random.randn(100))

in_oneD = np.array(in_oneD)
in_twoD = np.array(in_twoD)
target = np.array(target)

# データを学習用とテスト用に分割
oneD_train, oneD_test, twoD_train, twoD_test, target_train, target_test = train_test_split(in_oneD, in_twoD, target, test_size=1/7, random_state=0)

# PyTorchのテンソルに変換
oneD_train = torch.Tensor(oneD_train)
oneD_test = torch.Tensor(oneD_test)
twoD_train = torch.Tensor(twoD_train)
twoD_test = torch.Tensor(twoD_test)
target_train = torch.Tensor(target_train)
target_test = torch.Tensor(target_test)

# 入力と出力を結合
ds_train = TensorDataset(oneD_train, twoD_train, target_train)
ds_test = TensorDataset(oneD_test, twoD_test, target_test)

# DataLoaderを作成
train_loader = DataLoader(ds_train, batch_size=8, shuffle=True)
test_loader = DataLoader(ds_test, batch_size=1, shuffle=False)

ここでは、データセットを作成する際に、すべての入力と出力をTensorDatasetの引数として渡します。最終的にDataLoaderを作成する際には、このデータセットをそのまま渡します。

DataLoaderの使い方

DataLoaderは一般的に学習用関数内で使用します。以下はその実装例です：

# 学習用関数
def train(data_loader, model, optimizer, loss_fn, device):
    model.train()
    for in_oneD, in_twoD, target in data_loader:
        in_twoD = in_twoD.to(device, dtype=torch.float)
        in_oneD = in_oneD.to(device, dtype=torch.float)
        target = target.to(device, dtype=torch.float)

        optimizer.zero_grad()
        output = model(in_twoD, in_oneD)
        loss = loss_fn(output, target)

        loss.backward()
        optimizer.step()

    return loss.item()

ここで重要なのは、DataLoaderから取得したデータをモデルのforwardメソッドに渡す際の順序です。この順序は、モデルのforwardメソッドで定義した引数の順序に従います。

まとめ

PyTorchで複数の入力がある場合の実装方法について解説しました。以下に示すコードは、説明した内容を実際に動作させるための全体のソースコードです。PyTorchの環境が整っていれば、このコードを実行することで動作を確認できます。

#!/usr/bin/env python3
# -*- coding: UTF-8 -*-

import numpy as np
import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.block1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=1),
            nn.BatchNorm2d(32)
        )
        self.block2 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=1),
            nn.BatchNorm2d(64)
        )
        self.full_connection = nn.Sequential(
            nn.Linear(in_features=64*13*18+3, out_features=1024),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(in_features=1024, out_features=100, bias=False)
        )

    def forward(self, in_twoD, in_oneD):
        x = self.block1(in_twoD)
        x = self.block2(x)
        x = x.view(x.size(0), -1)
        x = torch.cat([x, in_oneD], dim=1)
        y = self.full_connection(x)
        return y

def train(data_loader, model, optimizer, loss_fn, device):
    model.train()
    for in_oneD, in_twoD, target in data_loader:
        in_twoD = in_twoD.to(device, dtype=torch.float)
        in_oneD = in_oneD.to(device, dtype=torch.float)
        target = target.to(device, dtype=torch.float)

        optimizer.zero_grad()
        output = model(in_twoD, in_oneD)
        loss = loss_fn(output, target)

        loss.backward()
        optimizer.step()

    return loss.item()

def main():
    # サンプルデータの作成
    in_oneD, in_twoD, target = [], [], []

    for _ in range(100):
        in_oneD.append(np.random.randn(3))
        in_twoD.append(np.random.randn(1, 15, 20))
        target.append(np.random.randn(100))

    in_oneD = np.array(in_oneD)
    in_twoD = np.array(in_twoD)
    target = np.array(target)

    # データを学習用とテスト用に分割
    oneD_train, oneD_test, twoD_train, twoD_test, target_train, target_test = train_test_split(
        in_oneD, in_twoD, target, test_size=1/7, random_state=0)

    # PyTorchのテンソルに変換
    oneD_train = torch.tensor(oneD_train, dtype=torch.float)
    oneD_test = torch.tensor(oneD_test, dtype=torch.float)
    twoD_train = torch.tensor(twoD_train, dtype=torch.float)
    twoD_test = torch.tensor(twoD_test, dtype=torch.float)
    target_train = torch.tensor(target_train, dtype=torch.float)
    target_test = torch.tensor(target_test, dtype=torch.float)

    # 入力と出力を結合
    ds_train = TensorDataset(oneD_train, twoD_train, target_train)
    ds_test = TensorDataset(oneD_test, twoD_test, target_test)

    # DataLoaderを作成
    train_loader = DataLoader(ds_train, batch_size=8, shuffle=True)
    test_loader = DataLoader(ds_test, batch_size=1, shuffle=False)

    # デバイス（GPU/CPU）の設定
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 初期モデルを作成
    model = Model().to(device)

    # 損失関数とOptimizerを作成
    loss_fn = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.01)

    # 学習
    for epoch in range(100):
        train_loss = train(train_loader, model, optimizer, loss_fn, device)
        print(f'Epoch {epoch+1}, Loss: {train_loss:.4f}')

if __name__ == '__main__':
    main()
スポンサーリンク