PyTorch のモデルを GPU に配置できない

Question

Pytorch を使ってモデルを構築しようとしています。最初の段階で、Backboneモデルから中間層を３つ取り出すモジュールを書いていたのですが、これを GPU 上に配置出来なくて困っています。

import torch.nn as nn
import torchvision.models as models


class Backbone(nn.Module):
    def __init__(self, backbone_type="ResNet50"):
        super(Backbone, self).__init__()
        self.backbone_type = backbone_type

    def forward(self, x):
        if self.backbone_type == "ResNet50":
            model = models._utils.IntermediateLayerGetter(
                models.resnet50(), {"layer1": 2, "layer2": 3, "layer3": 4}
            )

            return model(x)

テストコードとして以下のようなコードを書いています。

import numpy as np
import pytest
import torch

from src.model import Backbone

@pytest.mark.skipif(not torch.cuda.is_available(), reason="There is no GPU")
def test_run_backbone_on_gpu():
    back = Backbone()
    dummy_input = torch.from_numpy(np.random.random((1, 3, 224, 224))).to(
        "cuda", dtype=torch.float
    )
    back = back.cuda()
    dummy_output = back(dummy_input)
    assert dummy_output[2].size() == torch.Size([1, 256, 56, 56])
    assert dummy_output[3].size() == torch.Size([1, 512, 28, 28])
    assert dummy_output[4].size() == torch.Size([1, 1024, 14, 14])

このテストを実行すると以下のような Runtime Error が出ます。

================================================================================================ FAILURES =================================================================================================
________________________________________________________________________________________ test_run_backbone_on_gpu _________________________________________________________________________________________

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="There is no GPU")
    def test_run_backbone_on_gpu():
        back = Backbone()
        dummy_input = torch.from_numpy(np.random.random((1, 3, 224, 224))).to(
            "cuda", dtype=torch.float
        )
        back = back.cuda()
>       dummy_output = back(dummy_input)

tests/test_model.py:43:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
.venv/lib/python3.8/site-packages/torch/nn/modules/module.py:722: in _call_impl
    result = self.forward(*input, **kwargs)
src/model.py:16: in forward
    return model(x)
.venv/lib/python3.8/site-packages/torch/nn/modules/module.py:722: in _call_impl
    result = self.forward(*input, **kwargs)
.venv/lib/python3.8/site-packages/torchvision/models/_utils.py:63: in forward
    x = module(x)
.venv/lib/python3.8/site-packages/torch/nn/modules/module.py:722: in _call_impl
    result = self.forward(*input, **kwargs)
.venv/lib/python3.8/site-packages/torch/nn/modules/conv.py:419: in forward
    return self._conv_forward(input, self.weight)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
input = tensor([[[[0.1313, 0.0234, 0.7552,  ..., 0.3833, 0.4139, 0.8594],
          [0.1304, 0.7056, 0.5246,  ..., 0.4825, 0.6..., 0.3003, 0.3479, 0.5493],
          [0.3813, 0.9502, 0.2774,  ..., 0.6596, 0.5868, 0.9608]]]],
       device='cuda:0')
weight = Parameter containing:
tensor([[[[-4.6613e-02, -9.6919e-04,  8.5332e-03,  ...,  2.0174e-02,
            9.1100e-03, -4....[-2.4699e-02,  1.7219e-02, -1.9792e-02,  ..., -7.8135e-03,
           -3.0772e-02, -1.7609e-02]]]], requires_grad=True)

    def _conv_forward(self, input, weight):
        if self.padding_mode != 'zeros':
            return F.conv2d(F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode),
                            weight, self.bias, self.stride,
                            _pair(0), self.dilation, self.groups)
>       return F.conv2d(input, weight, self.bias, self.stride,
                        self.padding, self.dilation, self.groups)
E       RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

.venv/lib/python3.8/site-packages/torch/nn/modules/conv.py:415: RuntimeError
========================================================================================= short test summary info =========================================================================================
FAILED tests/test_model.py::test_run_backbone_on_gpu - RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

このエラーは、入力値は GPU 上にあるけど モデル（重みパラメータ）が GPU 上にないということを言っているのだと思いますが、これをどのように解決すればいいのかわかりません。

よろしくお願いいたします。

Comments

[meg_]

エラーメッセージは全文掲載してください。

[Yhaya]

追加しました。よろしくお願いします

Answer 1

解決方法

Backbone() 内で使用する model を forward() ではなく、 __init__() 内で属性として定義する。

class Backbone(nn.Module):
    def __init__(self, backbone_type="ResNet50"):
        super(Backbone, self).__init__()
        self.backbone_type = backbone_type
        if self.backbone_type == "ResNet50":
            self.model = models._utils.IntermediateLayerGetter(
                models.resnet50(), {"layer1": 2, "layer2": 3, "layer3": 4}
            )

    def forward(self, x):
        return self.model(x)

原因について

Pytorch ではモデルを構成する層などは nn.Module を継承しています。
複数の nn.Module を組み合わせてモデルを作るわけですが、nn.Module 内で使用する nn.Module は、__init__() 内で属性として定義しないと、その nn.Module の子として自動的に認識されません。

__init__() 内で属性として定義した場合

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # このモデルで使用する nn.Module はここで属性として定義する
        self.linear1 = torch.nn.Linear(10, 10)
        self.linear2 = torch.nn.Linear(10, 10)

    def forward(self, x):
        y = self.linear1(x)
        y = self.linear2(y)
        return y


net = Net()
# 子として認識されている
for child in net.children():
    print(child)
# Linear(in_features=10, out_features=10, bias=True)
# Linear(in_features=10, out_features=10, bias=True)

__init__() 内で属性として定義しなかった場合

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # このモデルで使用する nn.Module はここで属性として定義する
        self.linear2 = torch.nn.Linear(10, 10)

    def forward(self, x):
        linear1 = torch.nn.Linear(10, 10)
        y = linear1(x)
        y = self.linear2(y)
        return y


net = Net()
# linear1 が子として認識されていない
for child in net.children():
    print(child)
# Linear(in_features=10, out_features=10, bias=True)

to("cuda") や cuda() は nn.Module 内の子の nn.Module を GPU に転送するというものですが、そもそも子として認識されていない nn.Module は転送されないままということになります。
質問のコードだと models._utils.IntermediateLayerGetter() の返り値を __init__() 内で属性として定義していなかったので、子として認識されておらず、cuda() を実行しても GPU に転送されないままとなっています。

コードを読むと、nn.module 周りがどのように実現されているかがわかります。

torch.nn.modules.module — PyTorch 1.6.0 documentation

Comments

[Yhaya]

なるほど、わかりやすい解説ありがとうございます。勉強になりました。