Deep dreamが出た頃はまだtensorflowも公開されておらず、ディープラーニングはpythonというよりは、LuaやC++で書かれていた。懐かしのDeep dreamをpytorchを使って実装してみる。

下記のようなアナザーワールドを作ってみよう。

f:id:tzmi:20200310223348j:plain

実装上のポイント

実装上のポイントを挙げてみる

小さい画像から始めてだんだん大きくしていく
大きくしていく段階で解像度の低い画像を差し引いて解像度の高い画像を加えていく
modelのパラメータを止めて入力画像xを更新する
lossはシンプルに出力の二乗和
lossが大きくなる勾配方向にxを動かす、optimizerは使わずに自分で更新を書く

モデルの定義

VGG16のpretrained modelから途中の層までを抽出する。下記の例では27番目の層までを抽出している。さらにパラメータの勾配を更新しないように止めておく。

    # define model                                                                             
    model = models.vgg16(pretrained=True).features
    layers = list(model.children())
    model = nn.Sequential(*layers[:27]).to(device).eval()
    for p in model.parameters():
        p.requires_grad_(False)

dataの用意

データを用意してtransformsでテンソル化と正規化を行う。処理後のx_orgは1,3,512,512の形を持つテンソルになる。

    # data
    img = skio.imread('col_org.jpg')
    img = cv2.resize(img, (512,512), interpolation=cv2.INTER_AREA)

    # transforms                                                                               
    mean = torch.tensor([0.485, 0.456, 0.406])
    std = torch.tensor([0.229, 0.224, 0.225])
    totensor = transforms.ToTensor()
    normalize = transforms.Normalize(mean=mean, std=std)
    transform = transforms.Compose([totensor, normalize])
    x_org = transform(img).unsqueeze(0)

学習

基本的には小さい画像から始めて、各大きさでniter回学習のイテレーションを回し、どんどん大きくしていく。あるサイズでの学習が終了したらアップサンプリングして次のサイズに行くが、小さい画像から始めると、もとの高解像度のエッジがない状態となる。そこで、あるサイズの画像の学習が終了したら、そのサイズでの入力画像を差し引く(ループの最後にある引き算)。差し引いた画像をアップサンプリングして、次の解像度での入力画像を足す（最初のif文）。こうすることで、アップサンプリング後にも元画像のエッジを維持することができる。

    for ioct in range(noct):
        print(ioct)

        s = x_org.shape[2]//int(oct_scale**(noct-ioct-1))
        x_base = F.adaptive_avg_pool2d(x_org, s).to(device)

        if ioct == 0:
            x = x_base
        else:
            x = x_base + F.interpolate(x, s, mode='bicubic', align_corners=False)

        x = x.detach().requires_grad_(True)
        for i in range(niter):
            y = model.forward(x)
            loss = y.norm()
            loss.backward()
            print(i, loss.item())

            avg_grad = torch.abs(x.grad).mean().item()
            norm_lr = lr/avg_grad
            x.data += norm_lr * x.grad.detach() # 勾配と逆方向に移動                           
            for j in range(3):
                upper = (1-mean[j])/std[j]
                lower = -mean[j]/std[j]
                x.data[:,j] = torch.clamp(x.detach()[:,j], lower, upper)

            x.grad.detach().zero_()

        x.data -= x_base

学習で更新するのはモデルではなく、入力画像。入力画像はnn.Moduleの持つparameters()を持っていないため、optimizerは使えない。入力画像で初期化した入力用レイヤを自作することもできるが、今回は入力xを直接更新する方法をとる。そのままだとxは中間層となってしまい、２周目のループで勾配を計算できなくなる。これを避けるために学習ループ(niter)の前でxを初期化する、つまり一度データだけを取り出してrequires_grad_(True)で入力層に変換する。

loss計算と誤差逆伝搬(backward)を行った後は、勾配更新の部分を自分で書く。まずは計算した勾配の絶対値平均を計算してこの値でlrを正規化する。正規化後のlrを使って入力画像xのデータ部分x.data（勾配部分x.gradと分ける必要がある）を更新する。更新でははlossが大きくなる方向にxを動かす。最後に大きくなりすぎた値と小さくなりすぎた値をmeanとstdを使って補正する。入力画像の勾配リセット関数(model.zero_gradにあたるもの)はx.grad.detach().zero_()と書く。

出力用の整形

最後に最適化？？されたxを整形する。

    # 出力用の整形                                                                             
    x = (x+x_base).to('cpu')
    mean = mean.reshape(1,3,1,1).to('cpu')
    std = std.reshape(1,3,1,1).to('cpu')
    x = torch.clamp((x*std + mean)*255, 0, 255)
    x = x.detach().type(torch.uint8)
    x = x.squeeze().permute(1,2,0)