用Pytorch实现DeepDream

model.py

import numpy as np
import torch
from torch import nn
from torchvision import models

class CustomResnet18(nn.Module):
    def __init__(self, end_layer):
        super().__init__()
        self.resnet18 = models.resnet.resnet18(pretrained=True)
        self.end_layer = end_layer
        
    def forward(self, x):
        # end_layer range : 1, 2, 3, 4
        x = self.resnet18.conv1(x)
        x = self.resnet18.bn1(x)
        x = self.resnet18.relu(x)
        x = self.resnet18.maxpool(x)

        layers = [self.resnet18.layer1,
                  self.resnet18.layer2,
                  self.resnet18.layer3,
                  self.resnet18.layer4]
        for i in range(self.end_layer):
            x = layers[i](x)
        return x

if __name__ == '__main__':
    model = CustomResnet18(4)
    x = torch.Tensor(np.random.normal(size=(1, 3, 224, 224)))
    s = model(x)

DeepDream.py

from matplotlib import pyplot as plt
from PIL import Image
import numpy as np
from scipy import ndimage as nd
from model import CustomResnet18
import torch

if __name__ == '__main__':
    # load img
    img = np.array(Image.open(r'img.jpg')) / 255.
    
    # pretrained resnet18
    model = CustomResnet18(4).cuda()

    # hyper params
    base_lr = 1e0
    steps = 250

    # multi scales
    tmp_img = img.copy()
    tmp_img = np.transpose(tmp_img[np.newaxis, :, :, :], [0, 3, 1, 2])
    tmp_img = nd.zoom(tmp_img, (1, 1, 0.125, 0.125))
    for i in range(4):
        scale = 0.125 * 2**i
        # dreaming
        for step in range(steps):
            # shift img
            shift_x, shift_y = np.random.randint(-10, 11, size=(2,))
            tmp_img = np.roll(np.roll(tmp_img, shift_x, -1), shift_y, -2)
            img_var = torch.autograd.Variable(torch.Tensor(tmp_img)).cuda()
            img_var.requires_grad = True

            # forward & backward
            out = model(img_var)
            loss = out.norm() # DeepDream希望网络更加确信输入属于某个类，所以目标函数是输出的范数
            loss.backward()

            # update
            ratio = np.abs(np.array(img_var.grad.data.tolist())).mean()
            lr = base_lr / ratio
            img_var = img_var.data.add(img_var.grad.data * lr).requires_grad_(True)

            # shift back
            tmp_img = np.array(img_var.data.tolist())
            tmp_img = np.roll(np.roll(tmp_img, -shift_x, -1), -shift_y, -2)
            img_var = torch.autograd.Variable(torch.Tensor(tmp_img)).cuda()
            img_var.requires_grad = True
            
            # log
            print('Step : {}/{} , Scale : {}'.format(step+1, steps, scale))
            
        # scale up
        if i < 3:
            tmp_img = nd.zoom(tmp_img, (1, 1, 2, 2))
    
    # get result
    dream_img = np.transpose(np.array(img_var.data.tolist())[0], [1, 2, 0])

    # save result
    np.save('dream_img.npy', dream_img)
    
    # plot result
    dream_img = (dream_img - dream_img.min()) / (dream_img.max() - dream_img.min()) * 1.8
    dream_img = np.clip(dream_img, 0, 1)
    
    plt.subplot(1, 2, 1)
    plt.imshow(img)
    plt.title('Origin')

    plt.subplot(1, 2, 2)
    plt.imshow(dream_img)
    plt.title('Dream')

    plt.savefig('DeepDreamWithOrigin.png')
    plt.show()

效果图