PyTorch LightningのGPU訓練で「RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same」と怒られた場合

---

PyTorch LightningはGPUで訓練しても「.to(device)」のようなデバイス指定をしなくていいのが売りですが、「デバイスが異なる」と怒られてハマってしまったので、解決法を示します。

おこったこと

PyTorch Lightningは、PyTorchの訓練のような「.cuda()」や「.to(device)」のような、デバイス指定を明示的に与えなくてもいいのが売りです。（公式サイトより）

しかし、モデルを自分で実装して訓練したところ、PyTorchにありがちな「RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same」と怒られてしまい、ハマってしまいました。

詳細に書くと、

• torchvisionやtimmから読み込んだモデルはこれが起こらない
• ただ、自分でモデルを定義すると起こる

というものです。この理由を見ていきます。

環境

• PyTorch 1.11.0+cu113
• PyTorch Lightning 1.5.10

コード

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torchvision
import torchmetrics
import pytorch_lightning as pl

class TenLayersModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layers = []
        for i in range(3):
            for j in range(3):
                if i == 0 and j == 0:
                    in_ch = 3
                elif j == 0:
                    in_ch = 64 * (2**(i-1))
                else:
                    in_ch = 64*(2**i)
                out_ch = 64*(2**i)
                self.layers.append(nn.Conv2d(in_ch, out_ch, 3, padding=1))
                self.layers.append(nn.BatchNorm2d(out_ch))
                self.layers.append(nn.ReLU())
        self.layers.append(nn.AdaptiveAvgPool2d((1, 1)))
        self.fc = nn.Linear(256, 10)

        self.train_acc = torchmetrics.Accuracy()
        self.val_acc = torchmetrics.Accuracy()        

    def forward(self, inputs):
        x = inputs
        for l in self.layers:
            x = l(x)
        x = x.view(x.shape[0], 256)
        x = self.fc(x)
        return x

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.parameters(), 0.1, 0.9)
        scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [70, 90], gamma=0.1)
        return [optimizer], [scheduler]

    def training_step(self, train_batch, batch_idx):
        x, y_true = train_batch
        y_pred = self.forward(x)
        loss = F.cross_entropy(y_pred, y_true)
        y_pred_label = torch.argmax(y_pred, dim=-1)

        self.train_acc.update(y_pred_label, y_true)
        self.log("train_loss", loss, prog_bar=True, logger=True)
        return loss

    def validation_step(self, val_batch, batch_idx):
        x, y_true = val_batch
        y_pred = self.forward(x)
        loss = F.cross_entropy(y_pred, y_true)
        y_pred_label = torch.argmax(y_pred, dim=-1)

        self.val_acc.update(y_pred_label, y_true)
        self.log("val_loss", loss, prog_bar=False, logger=True)
        return loss

    def validation_epoch_end(self, outputs):
        self.log("val_acc", self.val_acc.compute(), prog_bar=True, logger=True)
        self.log("train_acc", self.val_acc.compute(), prog_bar=True, logger=True)
        self.train_acc.reset()
        self.val_acc.reset()

class MyDataModule(pl.LightningDataModule):
    def __init__(self):
        super().__init__()

    def prepare_data(self):
        self.train_dataset = torchvision.datasets.CIFAR10(
            "./data", train=True, download=True,
            transform=torchvision.transforms.ToTensor())
        self.val_dataset = torchvision.datasets.CIFAR10(
            "./data", train=False, download=True,
            transform=torchvision.transforms.ToTensor())

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=256, num_workers=4, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=256, num_workers=4, shuffle=False)

def main():
    model = TenLayersModel()
    cifar = MyDataModule()
    logger_csv = pl.loggers.CSVLogger("outputs", name="lightning_logs_csv")
    logger_tb = pl.loggers.TensorBoardLogger("outputs", name="lightning_logs_tb")
    trainer = pl.Trainer(gpus=[1], max_epochs=5)

    trainer.fit(model, cifar)

if __name__ == "__main__":
    main()

エラー

  File "e:\...\train.py", line 56, in validation_step
    y_pred = self.forward(x)
  File "e:\...\train.py", line 34, in forward
    x = l(x)
  File "C:\Users\...\Roaming\Python\Python37\site-packages\torch\nn\modules\module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\...\AppData\Roaming\Python\Python37\site-packages\torch\nn\modules\conv.py", line 447, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "C:\Users\...\AppData\Roaming\Python\Python37\site-packages\torch\nn\modules\conv.py", line 444, in _conv_forward
    self.padding, self.dilation, self.groups)
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

解決法：レイヤーをPythonのリストではなく、nn.ModuleListに変える

直接的な原因は、TenLayersModelのコンストラクタ内のこの行です。

        # 変更前
        self.layers = []

これをnn.ModuleListに変えるだけでエラー消えます。nn.ModuleListでもリスト同様appendで追加できます。

        # 変更後
        self.layers = nn.ModuleList()

Lightning Moduleの中で「.cuda()」のようなGPU用のモデルに変換する処理が自動的に走っているのですが、Pythonのリストでは対応していなく、nn.ModuleListにすると対応してくれるのが理由として考えられます。自動変換に対応しているコレクションを使えば、PyTorch Lightningでも明示的に.cuda()しなくてOKです。

解決したのでめでたしめでたし。ただ初見でこのエラー出たときはドキッととしました。

---

---

Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内

---