Pytorch GAN discriminatorの出力を入力サイズによって変更したい

Question

前提・実現したいこと

Discriminatorは、入力サイズによって、層を追加したいのですが、forward内にif分を書きまくるということになっています。
できるだけ、簡潔に書きたいのですが、教えて頂けると幸いです。

該当のソースコード

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()

        self.fromRGBs = nn.Conv2d(3,512,1,1,0)
        self.layer1 = self.layer(512,512,4,2,1)
        self.layer2 = self.layer(512,256,4,2,1)
        self.layer3 = self.layer(256,128,4,2,1)

    def layer(self, in_channel, out_channel, kernel_size, stride, padding):
        layer = nn.Sequential(
            nn.Conv2d(
                in_channels=in_channel,
                out_channels=out_channel,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
            ),
            nn.BatchNorm2d(out_channel),
            nn.LeakyReLU(inplace=True)
        )
        return layer

    def forward(self, x):
        from_rgb= self.fromRGBs(x)
        out = self.layer1(from_rgb)

        if x.shape[2]==8:
            out = self.layer2(out)
            return out

        elif x.shape[2]==16:
            out = self.layer2(out)
            out = self.layer3(out)
            return out

        return out

if __name__ == '__main__':
    d = Discriminator()
    image = torch.randn(size=(10,3,16,16))
    out = d(image)
    print(out.shape)

Answer 1

そのソースでは分岐が2個しかないですが、本来はさらにelifが大量に必要だけど追加したくない、という事でしょうか。

このソースから読み解けるルール

「x.shape[2]は8の倍数で、8増えるごとにレイヤーをかぶせる」

で考えた案です。

class TestDiscriminator:
    def __init__(self):
        self.fromRGBs = lambda x: "rgb" + str(x.shape[2])  # テスト用
        self.layers = [   # ★★ layerをListに持つ
            self.layer(1),
            self.layer(2),
            self.layer(3)
        ]

    def layer(self, l_no):  # テスト用
        def lyr(from_rgb):
            return f"layer{l_no}_{from_rgb}"
        return lyr

    def forward(self, x):

        # ★★ shape_2を8ずつ減らして0まで繰り返す。
        def recursive(dep, shape_2, ob):
            ret = self.layers[dep](ob)
            return ret if shape_2 == 0 else recursive(dep + 1, shape_2 - 8, ret)

        return recursive(0, x.shape[2], self.fromRGBs(x))  # ★★再帰でループ


if __name__ == "__main__":
    class S:  # テスト用
        def __init__(self, i):
            self.shape = [0, 0, i]

    d = TestDiscriminator()

    print(d.forward(S(0)))   # x.shape[2] = 0
    print(d.forward(S(8)))   # x.shape[2] = 8
    print(d.forward(S(16)))  # x.shape[2] = 16

# 実行結果
# layer1_rgb0
# layer2_layer1_rgb8
# layer3_layer2_layer1_rgb16

「#テスト用」と書いたところはPytorchなしで単体で動かすために改造したところなので、無視して下さい。

※Pytorch動かせる環境を持っておらず、機械学習の知識も無いので、「素のPythonコードとして」のみで考えてます。筋違いならすいません。

ルールが実は単純でない、または分岐は2個から増えるわけではないという場合、あまり小細工すると余計に複雑になってメンテ難しくなるので、愚直にif文でも仕方ないかもしれません。