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

この記事では、PyTorchで複数の入力を受け取る深層学習モデルを実装する方法を解説します。モデルの定義、DataLoaderの作り方と使い方の順に説明し、最後にそのまま実行できる全体のコードを示します。

掲載しているコードはPyTorch 2.x系を想定しています。デバイスの選択にtorch.acceleratorを使っているため、その部分だけはPyTorch 2.6以降が必要です。それより前のバージョンを使っている場合の書き方も本文中に添えます。

目次

  1. モデルの実装方法
  2. DataLoaderの作り方
  3. DataLoaderの使い方
  4. まとめ

1. モデルの実装方法

入力が2か所あるネットワークを実装します。畳み込み層には「チャネル1×高さ15×幅20」のデータを入力し、その出力に「3成分のベクトル」をつないでから全結合層へ渡す、という構成です。

ネットワークのコードは次の通りです。

import torch
from torch import nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        # 2次元データを処理する畳み込みブロック
        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),
        )
        # 畳み込み側の出力64×13×18に、ベクトル入力の3を足した数が入力数になる
        self.full_connection = nn.Sequential(
            nn.Linear(in_features=64 * 13 * 18 + 3,
                      out_features=1024),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=1024,
                      out_features=100,
                      bias=False),
        )

    def forward(self, in_twoD, in_oneD):
        # 2次元データを畳み込みブロックに通す
        x = self.block1(in_twoD)
        x = self.block2(x)
        # (バッチ, 64, 13, 18) を (バッチ, 14976) に平坦化する
        x = torch.flatten(x, start_dim=1)
        # 平坦化した出力とベクトル入力を、特徴量の方向につなげる
        x = torch.cat([x, in_oneD], dim=1)
        return self.full_connection(x)

複数入力の要点はforwardにあります。畳み込み側の出力を平坦化し、torch.catでベクトル入力とつないでから全結合層へ渡します。連結する軸をdim=1にしているのは、バッチ方向ではなく特徴量の方向でつなぎたいからです。dim=0はバッチ方向なので、このコードで指定すると特徴量の数(14976と3)が合わずにエラーになります。

初心者がまずつまずくのは、全結合層のin_featuresに書く64 * 13 * 18 + 3という数字だと思います。これは自分で計算する必要があるので、内訳を確認しておきます。

  1. 入力は「チャネル1×高さ15×幅20」。
  2. Conv2dkernel_size=3, padding=1なので、縦横の大きさは変わりません。block1の畳み込みを通った時点で「チャネル32×15×20」。
  3. block1のMaxPool2d(kernel_size=2, stride=1)で縦横が1ずつ減り、「32×14×19」。
  4. block2の畳み込みでチャネルが64になり、プーリングでまた縦横が1ずつ減って「64×13×18」。
  5. 平坦化すると64 × 13 × 18 = 14976。ここにベクトル入力の3成分を足して14979が全結合層の入力数です。

入力サイズやカーネルサイズを変えると、この数字も変わります。手で計算し直すのが面倒なら、nn.Linearの代わりに入力数を自動で決めてくれるnn.LazyLinear(1024)を使う方法もあります。ただしその場合、パラメータが確定するのは最初に一度forwardを実行した後なので、それより前にoptim.Adam(model.parameters(), ...)を作るとエラーになります。まずは自分で計算して、仕組みを掴んでからのほうが分かりやすいと思います。

なお古い解説記事では、super(Model, self).__init__()や、平坦化にx.view(x.size(0), -1)を使う書き方をよく見かけます。前者はPython 3ならsuper().__init__()で足ります。後者は、テンソルがメモリ上で連続していない場合にエラーになることがあるので、必要なら自動でコピーを作るtorch.flattenのほうが安全です。

2. DataLoaderの作り方

次にDataLoaderを作ります。複数の入力を扱う部分に集中したいので、データの中身は乱数で生成します。

import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset, random_split

def build_dataloaders(samples=100, batch_size=8):
    # サンプルデータの作成(実際に使うときは、ここを自分のデータに置き換える)
    rng = np.random.default_rng(0)
    in_twoD = rng.standard_normal((samples, 1, 15, 20), dtype=np.float32)
    in_oneD = rng.standard_normal((samples, 3), dtype=np.float32)
    target = rng.standard_normal((samples, 100), dtype=np.float32)

    # 入力と正解をまとめる。モデルのforwardと同じ順序で渡すと取り違えにくい
    dataset = TensorDataset(
        torch.from_numpy(in_twoD),
        torch.from_numpy(in_oneD),
        torch.from_numpy(target),
    )

    # 学習用とテスト用に分割する(テスト用は全体の約7分の1)
    n_test = len(dataset) // 7
    ds_train, ds_test = random_split(
        dataset,
        [len(dataset) - n_test, n_test],
        generator=torch.Generator().manual_seed(0),  # 分け方を毎回同じにする
    )

    # DataLoaderを作成する
    train_loader = DataLoader(ds_train, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(ds_test, batch_size=batch_size, shuffle=False)
    return train_loader, test_loader

複数入力で大事なのはTensorDatasetの部分です。入力が何個あっても、すべてのテンソルを引数として並べて渡すだけで、DataLoaderは同じ順番のタプルを返してくれます。入力が増えるほど順序を間違えやすくなるので、モデルのforwardの引数と同じ並びにしておくと安全です。ここではforward(in_twoD, in_oneD)に合わせて、2次元データ、ベクトル、正解の順にしています。

細かい点ですが、NumPy配列からテンソルを作るときはtorch.from_numpyを使ってください。古い記事でよく見るtorch.Tensor(array)は型が明示されず、意図しない変換になることがあります。NumPy側でfloat32にしておけば、そのまま学習に使える型のテンソルになります。乱数もnp.random.randnではなく、np.random.default_rngで作ったGeneratorを使うのが現在の推奨です。

学習用とテスト用の分割にはtorch.utils.data.random_splitを使いました。scikit-learnのtrain_test_splitでも分割できますが、入力の数だけ戻り値が増えて読みにくくなります。データセットを作ってから分ければ、入力が何個あっても書き方は変わりません。

3. DataLoaderの使い方

DataLoaderは学習用と評価用の関数の中で回します。

def train_one_epoch(data_loader, model, optimizer, loss_fn, device):
    model.train()  # DropoutとBatchNormを学習用の動作に切り替える
    total_loss = 0.0
    for in_twoD, in_oneD, target in data_loader:
        # 取り出す順序は、TensorDatasetに渡した順序と同じ
        in_twoD = in_twoD.to(device)
        in_oneD = in_oneD.to(device)
        target = target.to(device)

        optimizer.zero_grad()             # 前回の勾配を消す
        output = model(in_twoD, in_oneD)  # 渡す順序はforwardの引数定義に合わせる
        loss = loss_fn(output, target)
        loss.backward()                   # 勾配を計算する
        optimizer.step()                  # 重みを更新する

        total_loss += loss.item()

    # ミニバッチごとの損失を平均して、このエポックの損失とする
    return total_loss / len(data_loader)

def evaluate(data_loader, model, loss_fn, device):
    model.eval()  # DropoutとBatchNormを推論用の動作に切り替える
    total_loss = 0.0
    with torch.inference_mode():  # 勾配を計算しないので速く、メモリも節約できる
        for in_twoD, in_oneD, target in data_loader:
            in_twoD = in_twoD.to(device)
            in_oneD = in_oneD.to(device)
            target = target.to(device)

            output = model(in_twoD, in_oneD)
            total_loss += loss_fn(output, target).item()

    return total_loss / len(data_loader)

forで受け取る変数の順序はDataLoaderが返すタプルの順序、つまりTensorDatasetに渡した順序です。一方、モデルへ渡す順序はforwardの引数定義に従います。この2つが食い違っていても、形が偶然合っているとエラーにならず、入れ替わったデータで学習し続けてしまいます。複数入力の実装で一番気をつけたい箇所です。

評価用の関数を分けているのは、学習中の損失だけを見ていても、モデルが学習データを丸暗記しているだけなのか判断できないからです。model.eval()でDropoutを止め、torch.inference_mode()の中で計算します。書籍やチュートリアルでよく見るtorch.no_grad()でも動きますが、推論だけならinference_modeのほうが速くなります。

4. まとめ

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

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

import numpy as np
import torch
from torch import nn, optim
from torch.utils.data import DataLoader, TensorDataset, random_split

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        # 2次元データを処理する畳み込みブロック
        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),
        )
        # 畳み込み側の出力64×13×18に、ベクトル入力の3を足した数が入力数になる
        self.full_connection = nn.Sequential(
            nn.Linear(in_features=64 * 13 * 18 + 3,
                      out_features=1024),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=1024,
                      out_features=100,
                      bias=False),
        )

    def forward(self, in_twoD, in_oneD):
        # 2次元データを畳み込みブロックに通す
        x = self.block1(in_twoD)
        x = self.block2(x)
        # (バッチ, 64, 13, 18) を (バッチ, 14976) に平坦化する
        x = torch.flatten(x, start_dim=1)
        # 平坦化した出力とベクトル入力を、特徴量の方向につなげる
        x = torch.cat([x, in_oneD], dim=1)
        return self.full_connection(x)

def build_dataloaders(samples=100, batch_size=8):
    # サンプルデータの作成(実際に使うときは、ここを自分のデータに置き換える)
    rng = np.random.default_rng(0)
    in_twoD = rng.standard_normal((samples, 1, 15, 20), dtype=np.float32)
    in_oneD = rng.standard_normal((samples, 3), dtype=np.float32)
    target = rng.standard_normal((samples, 100), dtype=np.float32)

    # 入力と正解をまとめる。モデルのforwardと同じ順序で渡すと取り違えにくい
    dataset = TensorDataset(
        torch.from_numpy(in_twoD),
        torch.from_numpy(in_oneD),
        torch.from_numpy(target),
    )

    # 学習用とテスト用に分割する(テスト用は全体の約7分の1)
    n_test = len(dataset) // 7
    ds_train, ds_test = random_split(
        dataset,
        [len(dataset) - n_test, n_test],
        generator=torch.Generator().manual_seed(0),  # 分け方を毎回同じにする
    )

    # DataLoaderを作成する
    train_loader = DataLoader(ds_train, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(ds_test, batch_size=batch_size, shuffle=False)
    return train_loader, test_loader

def train_one_epoch(data_loader, model, optimizer, loss_fn, device):
    model.train()  # DropoutとBatchNormを学習用の動作に切り替える
    total_loss = 0.0
    for in_twoD, in_oneD, target in data_loader:
        # 取り出す順序は、TensorDatasetに渡した順序と同じ
        in_twoD = in_twoD.to(device)
        in_oneD = in_oneD.to(device)
        target = target.to(device)

        optimizer.zero_grad()             # 前回の勾配を消す
        output = model(in_twoD, in_oneD)  # 渡す順序はforwardの引数定義に合わせる
        loss = loss_fn(output, target)
        loss.backward()                   # 勾配を計算する
        optimizer.step()                  # 重みを更新する

        total_loss += loss.item()

    # ミニバッチごとの損失を平均して、このエポックの損失とする
    return total_loss / len(data_loader)

def evaluate(data_loader, model, loss_fn, device):
    model.eval()  # DropoutとBatchNormを推論用の動作に切り替える
    total_loss = 0.0
    with torch.inference_mode():  # 勾配を計算しないので速く、メモリも節約できる
        for in_twoD, in_oneD, target in data_loader:
            in_twoD = in_twoD.to(device)
            in_oneD = in_oneD.to(device)
            target = target.to(device)

            output = model(in_twoD, in_oneD)
            total_loss += loss_fn(output, target).item()

    return total_loss / len(data_loader)

def main():
    torch.manual_seed(0)  # 実行するたびに結果が変わらないようにする

    # 使えるアクセラレータ(GPUなど)があればそれを使う。PyTorch 2.6以降の書き方
    device = (
        torch.accelerator.current_accelerator().type
        if torch.accelerator.is_available()
        else "cpu"
    )
    print(f"device: {device}")

    train_loader, test_loader = build_dataloaders()

    # モデル、損失関数、Optimizerを用意する
    model = Model().to(device)
    loss_fn = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)

    # 学習
    for epoch in range(1, 101):
        train_loss = train_one_epoch(train_loader, model, optimizer, loss_fn, device)
        if epoch % 10 == 0:
            test_loss = evaluate(test_loader, model, loss_fn, device)
            print(
                f"epoch {epoch:3d}  "
                f"train loss: {train_loss:.4f}  test loss: {test_loss:.4f}"
            )

if __name__ == '__main__':
    main()

デバイスの選択に使っているtorch.acceleratorはPyTorch 2.6で追加されたもので、CUDAだけでなくMPSやXPUも同じ書き方で扱えます(torch.accelerator)。2.6より前のバージョンでは、従来通りdevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")と書いてください。

おすすめ書籍

PyTorchでの深層学習モデルの実装をこれから学んでいきたい人は「つくりながら学ぶ! PyTorchによる発展ディープラーニング」がおすすめです。

(進化の早いこの界隈では)若干古い書籍ですが、古いおかげで一通りのモデルをゼロから実装する方法が書かれていますので、これを写経しながら読めば基本的な実装イメージは掴めるかなと思います。

入力を複数受け取るモデルの例としては、物体検出の基礎:アルゴリズムと応用事例深層学習を用いた3次元姿勢推定モデルとその応用事例で扱っているような構成が分かりやすいと思います。