PyTorchでGANの訓練をするときにrequires_grad（trainable）の変更はいるのかどうか

Posted On 2019-08-12

7k{icon} {views}

PyTorchでGANのある実装を見ていたときに、requires_gradの変更している実装を見たことがあります。Kerasだとtrainableの明示的な変更はいるんで、もしかしてPyTorchでもいるんじゃないかな？と疑問になったので、確かめてみました。

requires_gradの変更とは

あるレイヤーの係数を訓練するかどうかのフラグ。modelという変数があったときに、

for p in model.paramters():
    p.required_grad = False

とすることでそのモデル全体の係数を固定することができます。転移学習などに便利でしょう。

ものすごく簡単なGAN

検証用にものすごい簡単なGANのモデルを作ってみました。

import torch
from torch import nn

def weight_init(layer):
    if type(layer) in [nn.ConvTranspose2d, nn.Conv2d, nn.BatchNorm2d]:
        nn.init.constant_(layer.weight, 0.01)
        nn.init.zeros_(layer.bias)

class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Sequential(
            nn.ConvTranspose2d(100, 64, 6, 1), #1x1 -> 6x6
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.Conv2d(64, 3, 3, padding=1),
            nn.Tanh()
        )
        self.weights.apply(weight_init)

    def forward(self, inputs):
        return self.weights(inputs)

class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Sequential(
            nn.Conv2d(3, 32, 3, 2, padding=1), #6x6 -> 3x3
            nn.BatchNorm2d(32),
            nn.ReLU(True),
            nn.AvgPool2d(3), # 3x3 -> 1x1
            nn.Conv2d(32, 1, 1),
            nn.Sigmoid()
        )
        self.weights.apply(weight_init)

    def forward(self, inputs):
        return self.weights(inputs).view(inputs.size(0),-1)

CPUでも訓練できるぐらい簡単なモデルです。初期化はDともGともWeightを0.01の定数、biasを0としています。

初期値チェック

では初期値をチェックしてみましょう。

def initial_checking():
    model = Generator()
    print("--- Generator ---")
    for p in model.parameters():
        print(p)
    model = Discriminator()
    print("--- Disctiminator ---")
    for p in model.parameters():
        print(p)

これでGとDの係数一覧がでます。

--- Generator ---
Parameter containing:
tensor([[[[0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100],
          [0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100],
          [0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100],
          [0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100],
          [0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100],
          [0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100]],

## 中略
--- Disctiminator ---
Parameter containing:
tensor([[[[0.0100, 0.0100, 0.0100],
          [0.0100, 0.0100, 0.0100],
          [0.0100, 0.0100, 0.0100]],

         [[0.0100, 0.0100, 0.0100],
          [0.0100, 0.0100, 0.0100],
          [0.0100, 0.0100, 0.0100]],

DもGも0.01で初期化されているのが確認できました。ここには書いていませんがbiasは0で初期化されています。

Gの訓練を想定

よくあるGの訓練の1ステップを取り出します。本物のデータは適当に作っています。

def train_default():
    model_G = Generator()
    model_D = Discriminator()
    param_G = torch.optim.Adam(model_G.parameters())
    param_D = torch.optim.Adam(model_D.parameters())

    # 本物画像を適当におく
    real_img = torch.arange(0.0, 1.0, 1 / 6).view(6, 1, 1, 1).expand(6, 100, 1, 1)

    # Dのrequire_gradの変更なしにGをアップデート
    torch.manual_seed(123)
    fake_img = model_G(torch.randn(6, 100, 1, 1))
    out = model_D(fake_img)
    loss = nn.BCELoss()(out, torch.ones(6, 1))  # 本物として訓練
    print(loss)

    # backprop
    model_D.zero_grad()
    model_G.zero_grad()
    loss.backward()
    param_G.step()

    # Dのパラメーターチェック
    for p in model_D.parameters():
        print(p)

ここではrequires_gradの変更は一切行っていません。param_G.step()でGenerator側の更新は行っていますが、Dの係数は変更されるでしょうか？　2つの可能性が考えられます。

param_G.step()でGenerator側だけ更新してるんだから、D側の係数は訓練されるわけがない
いやいや、Gのロス計算しているときにDのモデルを使ってるんだから、Gを更新したときにDまで更新されちゃうでしょ

ちなみにKerasだと2で、G/Dの更新で明示的にtrainable(PyTorchでのrequires_grad)を変更する必要があります。PyTorchではどちらでしょうか？

答えは以下の通りです。

# 前略
Parameter containing:
tensor([0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100], requires_grad=True)
Parameter containing:
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0.], requires_grad=True)
# 後略

答えは1が正解。param_G.step()としているのでDは更新されません。ちなみにPyTorchの公式のDCGAN例ではこちらの方法を使っています。

params_D.step()とすると…

ではちょっと意地悪して、params_G.step()からparams_D.step()に変えてみます。Dは更新されるでしょうか？

Parameter containing:
tensor([0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110,
        0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110,
        0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110, 0.0110,
        0.0110, 0.0110, 0.0110, 0.0110, 0.0110], requires_grad=True)
Parameter containing:
tensor([0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010,
        0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010,
        0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010, 0.0010,
        0.0010, 0.0010, 0.0010, 0.0010, 0.0010], requires_grad=True)

結果はDが更新されました。もとの初期値は0.01と0だったのを思い出しましょう。

もうちょっと意地悪をする

もう少し意地悪をして、params_D.step()の手前にmodel_D.zero_grad()を入れてみます。

    # backprop
    model_D.zero_grad()
    model_G.zero_grad()
    loss.backward()
    model_D.zero_grad()
    param_D.step()

こんな感じ。どうなるでしょう？

Parameter containing:
tensor([0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100], requires_grad=True)
Parameter containing:
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0.], requires_grad=True)

当然更新されません。zero_grad()とは文字通り勾配を初期化しなさいという意味だったんですね。

params_D.step()からGの更新を見る

本来こういうケースはありませんが、D.step()して、Gの係数が変わっているか見てみましょう。こんな感じ。

    # backprop
    model_D.zero_grad()
    model_G.zero_grad()
    loss.backward()
    param_D.step()

    # Dのパラメーターチェック
    for p in model_G.parameters():
        print(p)

Parameter containing:
tensor([0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100, 0.0100,
        0.0100], requires_grad=True)
Parameter containing:
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],

結果：Gの係数は変わっていません。params_G.step()のときのDと同様係数は変わりませんでした。

結論

PyTorchにおいてGANを訓練するときは、Gの訓練をするときにreqires_grad=Falseにする必要はない。GとDで明示的にrequires_gradの切り替えはいらない。公式実装は正しい。ということでした。

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

技術書コーナー

北海道の駅巡りコーナー

Tags:DeepLearning, GAN, 小ネタ