前回はSAGANでやったので今回はPGGANを実装したいと思います．
実装したものはここ

環境

• python 3.8.0
• pytorch 1.9.0
• GPU GTX1080

どんな工夫がされているの？

Progressive Growing

生成したい画像のまま学習するのではなく段階的に学習させていきます．
論文中では4×4のサイズから学習を初めており．80k枚の画像を使用してサイズを遷移させ，その後80k枚の画像を使用して遷移後のサイズでの学習を行います．(今回は1サイズの半分のエポックで遷移させ残りのエポックで学習をさせました．)．

Pixelwise Normalization

PGGANの論文[1]によるとBatch Normalizationの代わりにPixelwise Normalizationを使用しGeneratorの3×3の畳み込み層の後で使用するとのことです．また式は下記のようになります．
b_{x,y}=\frac{a_{x,y}}{\sqrt{\frac{1}{N}\sum_{j=0}^{N-1}(a_{x,y}^j)^2+\epsilon}}　(\epsilon=10^{-8})
x,yは画像ピクセルの座標であり，変数jは各特徴量マップを表します．
式の通りに実装するとこんなかんじ？

class PixelNorm2d(nn.Module):
    def __init__(self,epsilon=1e-8):
        super().__init__()
        self.epsilon = epsilon

    def forward(self,inputs):
        denominator = torch.rsqrt(torch.mean(inputs**2,dim=1,dtype=inputs.dtype,keepdim=True) + self.epsilon)
        output = inputs * denominator
        return output

Equalized Leaning Rate

重みwをN(0,1)で初期化その後各層が実行される度にw=\frac{w}{c}の演算を行います

class EqualizedLRConv2d(nn.Conv2d):
    def __init__(self,in_channels,out_channels,kernel_size,stride=1,padding=0,bias=True,**kwargs):
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            bias=bias,
            **kwargs
            ) 
        nn.init.normal_(self.weight,mean=0,std=1)
        nn.init.constant_(self.bias,val=0.0)
        if(not isinstance(kernel_size,tuple)):
            kernel_size = (kernel_size,kernel_size)
        f_in = torch.prod(torch.tensor([in_channels,*kernel_size],dtype=self.weight.dtype,device=self.weight.device))
        self.scale_factor = torch.sqrt(2/f_in)
    def forward(self,inputs):
        inputs = self.scale_factor *inputs
        output = F.conv2d(inputs,self.weight,self.bias,self.stride,self.padding,self.dilation,self.groups)
        return output

上のself.scale_factorがcに該当します．また畳み込み演算は1次変換であるため，重みではなく入力値との演算に変更してあります．

Minibatch Standard Deviation

ミニバッチ内の特徴量マップすべての標準偏差を計算し平均をとったスカラーを入力値の特徴量に付け加えます．(言葉で説明するより実装を見たほうが早いと思う)

class MiniBatchStddev(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self,inputs):
        b,_,h,w = inputs.shape
        std = torch.std(inputs,unbiased=False,dim=0)
        v = torch.mean(std)
        output = torch.cat((inputs,torch.full(size=(b,1,h,w),fill_value=v.item(),dtype=inputs.dtype,device=inputs.device)),dim=1)
        return output

Wasserstein loss with gradient penalty

工夫ではないですが一応損失関数の実装も載せておきます．

class WassersteinGP(nn.Module):
    def __init__(self,net,penalty_coef=10):
        super().__init__()
        self.penalty_coef = penalty_coef
        self.net = net

    def forward(self,real,fake,*imgs):
        real_img,fake_img = imgs
        coef = torch.tensor(
            np.random.uniform(size=(real.shape[0],1,1,1)),
            requires_grad=True,
            device=real.device,
            dtype=real.dtype
            )
        penalty_input = fake_img * coef + (1 - coef)*real_img
        penalty_output = self.net(penalty_input)
        gradient = torch.autograd.grad(penalty_output,penalty_input,create_graph=True)[0]
        penalty = torch.square(torch.norm(gradient,dim=1)-1)
        output = fake.mean() - real.mean() + self.penalty_coef * penalty.mean()
        return output

学習結果

今回はSAGANのときと同じデータセットを使用して256×256の画像を生成しました．
バッチサイズ:4~64 16 ,128:8，256:5
画像サイズが小さいので16×16のサイズから載せます．
16×16 エポック0~5

32×32 エポック0~5

64×64 エポック0~5

128×128 エポック0~5

256×256ファイルサイズ大きくてグリットされたやつ貼れなかった（泣）

失敗例

感想

128×128までのサイズはカラーバリエーションもよく比較的に生成された画像を見ていて面白かった．しかし，256×256のサイズはバッチサイズが小さかったためか同じヘアカラーの画像が多く生成されるようになってしまった(原因は定かではないが多分バッチサイズ)．それに画像のノイズが他のサイズと比べて目立つ気がする．しかしそれを除けばわりかし大丈夫そう．

終わりに

今回は，PGGANを自分で実装してみました．今回実装するにあたって，クラス設計の大事さを学ぶことができました(GANとは関係ないけど)．Githubに上がっているコード(Pytorchの実装とか)を読んでいると自分のクラス設計の汚さに気が付きます(オブジェクト指向勉強しなきゃ…)．基本的に人に書いたコードを見せる機会がないので煩雑なコードを書きがちになってしまっている気がする．部内の人にコードレビューしてもらうのもいいかもしれないですね．次は発展型であるStyleGANやStyleGAN2，LightWeightGANを実装していきたいと思います．

参考

[1]Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. 2017.
Progressive Growing of GANs for Improved Quality, Stability, and Variation. arXiv preprint arXiv:1710.10196