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

PyTorchでDeep Learningを実装する際に，データを入力する箇所がネットワーク内に複数ある場合の実装方法についてです。

①モデル／②DataLoaderの作り方／③DataLoaderの使い方を順に確認していった後に，最後にまとめて全体のソースを記載しておきます。

なお，PyTorchでの基本的な実装方法は【PyTorch】MNISTのサンプルを動かしてみたを参考にしてみて下さい。

モデルの実装方法

今回はサンプルとして，下図のような入力が2か所あるネットワークを実装してみます。

まず最初に畳み込み層にINPUT(1)：”高さ15×幅20×チャネル1″の2次元行列を入力します。その後，畳み込み層からの出力とINPUT(2)：”3成分のベクトル”を一緒に全結合層に入力するような例です。

まずネットワークの書き方のソースは以下の通りです。

import torch
import torch.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)

        # 畳み込み層からの出力を1次元化
        x = x.view(x.size(0), 64*13*18)

        # 1次元化した畳み込み層のからの出力と2つめの入力を結合
        x = torch.cat([x, in_oneD], dim=1)

        # 全結合層に入力して計算
        y = self.full_connection(x)
        
        return y

各層を定義するinit関数では，27行目がポイントになります。in_featuresの値を畳み込み層からの入力数（64×13×18つ）とINPUT(2)データからの入力数（3つ）を足した値を設定します。

次にforward関数の中では44行目の部分がポイントです。ここでは，畳み込み層からの入力をviewメソッドで1次元化したあとに，torch.catメソッドを使ってINPUT(2)のデータと結合して2つの入力をひとつにまとめた形式に変換します。

DataLoaderの作り方

次にDataLoaderの作り方です。今回は複数入力がある場合の実装方法に焦点を当てますので，データの値はランダムに適当に作成します。

import numpy as np

# サンプルデータの作成
#   - 入力
#     - in_oneD : 3成分のベクトル形式のデータ（データ数：100）
#     - in_twoD : 15行20列の行列形式のデータ（データ数：100）
#   - ラベル
#     - target    : 100成分のベクトル形式のデータ（データ数：100）    
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.LongTensor(target_train)
target_test = torch.LongTensor(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)

ここでは，32～33行目でデータセットをまとめて作成しているところがポイントになります。一度覚えてしまえば単純な話ではありますが，すべての入力データと出力データをまとめて引数に与えてデータセットを作成します。

最終的にDataLoaderを作成するときは，ここで作成したデータセットをそのまま引数として与えればいいだけですので特に変わったポイントはありません。

DataLoaderの使い方

DataLoaderはよくあるパターンとしてはtrain関数を用意して，その中で使用するという形が多いかと思います。そこでの実装方法は以下の通りです。

# 学習用関数
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()

ここでポイントになるのは7行目と14行目です。

まず7行目についてですが，DataLoaderからはデータセットを作成したときに与えた引数の形でデータがfeedされます。そのため，”for in_oneD, in_twoD, target in data_loader”という形で3つのデータを受け取れます。また，その受け取ったデータは通常と同じように，それぞれ個別にto関数でGPUに送ってやれば問題ありません。

次に14行目ですが，ここでは2つの入力のデータを引数としてモデルに与えてやります。引数の順番は，”モデルの実装方法”で作成したforward関数で指定した引数の順番と同じにすれば良いです。

まとめ

PyTorchで複数の入力がある場合の実装方法のポイントは以上の通りです。一度覚えてしまえばそれ程難しくないというか，そのまま実装しているだけだなというのが分かるかと思います。

最後に，そのまま学習を実行できる状態のソース全体を記載しておきます。PyTorchがインストール出来ている環境でこのソースを適当なPythonファイルとして保存して実行すれば動くはずです。

今回は複数の入力がある場合にポイントとなる箇所に焦点を当てた記事でしたので，損失関数やオプティマイザなどの細かい箇所の基本的なことが気になる場合は，【PyTorch】MNISTのサンプルを動かしてみたもご覧下さい。

#!/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), # '+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)

        # 畳み込み層からの出力を1次元化
        x = x.view(x.size(0), 64*13*18)

        # 1次元化した畳み込み層のからの出力と2つめの入力を結合
        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 : 3成分のベクトル形式のデータ（データ数：100）
    #     - in_twoD : 15行20列の行列形式のデータ（データ数：100）
    #   - ラベル
    #     - target    : 100成分のベクトル形式のデータ（データ数：100）    
    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.LongTensor(target_train)
    target_test = torch.LongTensor(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)

    # デバイス（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'Loss: {train_loss:.2f}')
        # 以下，省略


if __name__ == '__main__':
    main()