Pytorch学习笔记——风格迁移

import torch
import torch.nn as nn
from torch.autograd import Variable
import torchvision
from torchvision import transforms, models
from PIL import Image
import argparse
import numpy as np
import os

use_gpu = torch.cuda.is_available()# 是否能够使用GPU

dtype = torch.cuda.FloatTensor if use_gpu else torch.FloatTensor  #判断数据允许GPU 否则CPU tensor
#导入数据
def load_image(image_path, transforms=None, max_size=None, shape=None):
    image = Image.open(image_path)#打开图片  image_path 为图片路径
    image_size = image.size #图片大小

    if max_size is not None:
        #获取图像size，为sequence
        image_size = image.size
        #转化为float的array
        size = np.array(image_size).astype(float)
        size = max_size / size * size;
        image = image.resize(size.astype(int), Image.ANTIALIAS)

#Image.ANTIALIAS在当前的PIL版本中，这个滤波器只用于改变尺寸和缩略图方法。

if shape is not None: image = image.resize(shape, Image.LANCZOS)#Image.LANCZOS 缩小图片比例
#必须提供transform.ToTensor，转化为4D Tensor .unsqueeze(0)表示增维 .squeeze(0)为降维
if transforms is not None: image = transforms(image).unsqueeze(0) #是否拷贝到GPU return image.type(dtype)class VGGNet(nn.Module): def __init__(self): super(VGGNet, self).__init__() self.select = ['0', '5', '10', '19', '28']#设置需要提取的 self.vgg19 = models.vgg19(pretrained = True).features def forward(self, x): features = [] #name类型为str，x为Variable for name, layer in self.vgg19._modules.items(): x = layer(x) if name in self.select: features.append(x) return featuresdef main(config): #定义图像变换操作，必须定义.ToTensor()。 transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) #加速content和style图像，style图像resize成同样大小 content.size(0-3)：0表示batch 1表示通道，2，3表示图片大小 content = load_image(config.content, transform, max_size = config.max_size) style = load_image(config.style, transform, shape = [content.size(2), content.size(3)]) #.clone()将concent复制一份作为target，并需要计算梯度，作为最终的输出
target = Variable(content.clone(), requires_grad = True) optimizer = torch.optim.Adam([target], lr = config.lr, betas=[0.5, 0.999]) vgg = VGGNet() if use_gpu: vgg = vgg.cuda() for step in range(config.total_step): #分别计算5个特征图 target_features = vgg(target) content_features = vgg(Variable(content)) style_features = vgg(Variable(style)) content_loss = 0.0 style_loss = 0.0 for f1, f2, f3 in zip(target_features, content_features, style_features): #计算content_loss content_loss += torch.mean((f1 - f2)**2) # **表示上标 **2 表示平方 n, c, h, w = f1.size() #将特征reshape成二维矩阵相乘，求gram矩阵 f1 = f1.view(c, h * w) f3 = f3.view(c, h * w) f1 = torch.mm(f1, f1.t())#f1.t() 表示f1的转置，但不改变f1
f3 = torch.mm(f3, f3.t()) #计算style_loss style_loss += torch.mean((f1 - f3)**2) / (c * h * w) #计算总的loss loss = content_loss + style_loss * config.style_weight #反向求导与优化 optimizer.zero_grad() loss.backward() optimizer.step() if (step+1) % config.log_step == 0: print ('Step [%d/%d], Content Loss: %.4f, Style Loss: %.4f' %(step+1, config.total_step, content_loss.data[0], style_loss.data[0])) if (step+1) % config.sample_step == 0: # Save the generated image denorm = transforms.Normalize((-2.12, -2.04, -1.80), (4.37, 4.46, 4.44)) img = target.clone().cpu().squeeze() img = denorm(img.data).clamp_(0, 1) torchvision.utils.save_image(img, 'output-%d.png' %(step+1))if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('--content', type=str, default='content.jpg') parser.add_argument('--style', type=str, default='style.jpg') parser.add_argument('--max_size', type=int, default=400) parser.add_argument('--total_step', type=int, default=5000) parser.add_argument('--log_step', type=int, default=10) parser.add_argument('--sample_step', type=int, default=1000) parser.add_argument('--style_weight', type=float, default=100) parser.add_argument('--lr', type=float, default=0.003) config = parser.parse_args() print(config) main(config)

下图为Pytorch 自带VGG19的网络结构，程序所取卷积层为1，3，5，9，12。

左侧图片是content，右侧图片是style

下面是左上到右下分别是迭代1000，2000，3000，4000的风格迁移后的图片

Pytorch学习笔记——风格迁移

猜你喜欢