Pytorch笔记：风格迁移

Pytorch笔记：风格迁移
训练模型：风格迁移网络+VGG16网络
生成网络：风格迁移网络
代码如下（根据陈云《深度学习框架：Pytorch入门与实践》的代码改动）
main.py

import torch as t
import cv2 as cv
import torchvision as tv
from torch.utils import data
from transformer_net import TransformerNet
import utils
from PackedVGG import Vgg16
from torch.nn import functional as Fmean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]class Config(object):#基本参数配置# General Argsuse_gpu = Truemodel_path = None  # pretrain model path (for resume training or test)# Train Argsimage_size = 256  # image crop_size for trainingbatch_size = 3data_root = 'F:\dl\chapter7\data2/'  # dataset root：$data_root/coco/a.jpgnum_workers = 4  # dataloader num of workerslr = 1e-3epoches = 2  # total epoch to traincontent_weight = 1e5  # weight of content_lossstyle_weight = 1e10  # weight of style_lossstyle_path = 'F:\dl\chapter7\style.png'  # style image pathdebug_file = '/tmp/debugnn'  # touch $debug_fie to interrupt and enter# Test Argscontent_path = 'F:\dl\chapter7\input.jpg'  # input file to do style transfer [for test]result_path = 'F:\dl\output.png'  # style transfer result [for test]def train():opt = Config()device = t.device('cuda') if opt.use_gpu else t.device('cpu')#判断是否使用GPU训练# Data loadingtransfroms = tv.transforms.Compose([#8万多张平时场景图数据集tv.transforms.Resize(opt.image_size),tv.transforms.CenterCrop(opt.image_size),tv.transforms.ToTensor(),tv.transforms.Lambda(lambda x: x * 255)])dataset = tv.datasets.ImageFolder(opt.data_root, transfroms)dataloader = data.DataLoader(dataset, opt.batch_size)# style transformer network 风格迁移网络transformer = TransformerNet()if opt.model_path:#如果有训练一半的模型，加载后继续训练transformer.load_state_dict(t.load(opt.model_path, map_location=lambda _s, _: _s))transformer.to(device)# Vgg16 for Perceptual Lossvgg = Vgg16().eval()#训练和测试时采用不同方式 Dropout和Batch Normalizationvgg.to(device)for param in vgg.parameters():param.requires_grad = False# Optimizeroptimizer = t.optim.Adam(transformer.parameters(), opt.lr)#风格迁移网络优化器# Get style image  获得所需风格图片style = utils.get_style_data(opt.style_path)style = style.to(device)# gram matrix for style image 计算风格图片的风格矩阵with t.no_grad():features_style = vgg(style)gram_style = [utils.gram_matrix(y) for y in features_style]for epoch in range(opt.epoches):for ii, (x, _) in enumerate(dataloader):# Trainoptimizer.zero_grad()x = x.to(device)y = transformer(x)y = utils.normalize_batch(y)x = utils.normalize_batch(x)features_y = vgg(y)features_x = vgg(x)# content loss 内容损失 生成图片和原图越像越好content_loss = opt.content_weight * F.mse_loss(features_y.relu2_2, features_x.relu2_2)# style loss 风格损失 生成图片和风格图的风格越像越好style_loss = 0.for ft_y, gm_s in zip(features_y, gram_style):#'relu1_2', 'relu2_2', 'relu3_3', 'relu4_3'输出相加gram_y = utils.gram_matrix(ft_y)style_loss += F.mse_loss(gram_y, gm_s.expand_as(gram_y))style_loss *= opt.style_weighttotal_loss = content_loss + style_loss#两种损失叠加print(ii,"loss is %f" % total_loss.data.cpu().numpy())total_loss.backward()optimizer.step()if ii%500==0 :t.save(transformer.state_dict(), 'F:\dl\chapter7\check/style_%s.pth' % ii)# save checkpointt.save(transformer.state_dict(), 'F:\dl\chapter7\check/%s_style.pth' % epoch)def stylize():"""perform style transfer"""opt = Config()#device = t.device('cuda') if opt.use_gpu else t.device('cpu')device=t.device('cpu')# input image preprocess 输入需变换图片content_image = tv.datasets.folder.default_loader('F:\dl\mmexport1598515552726.jpg')content_transform = tv.transforms.Compose([tv.transforms.ToTensor(),tv.transforms.Lambda(lambda x: x.mul(255))])content_image = content_transform(content_image)content_image = content_image.unsqueeze(0).to(device).detach()# model setup 加载模型style_model = TransformerNet().eval()style_model.load_state_dict(t.load('F:\dl\chapter7\check\style_25000.pth', map_location=lambda _s, _: _s))style_model.to(device)# style transfer and save output 输出生成图片output = style_model(content_image)output_data = output.cpu().data[0]tv.utils.save_image(((output_data / 255)).clamp(min=0, max=1), 'F:\dl\outputmym.png' )print("输出并保存完毕")scr = cv.imread('F:\dl\outputmym.png' )cv.imshow("scr", scr)cv.waitKey(0)#stylize()#输出风格化图片
train()#训练

utils.py

import torch as t
import torchvision as tv
import numpy as npIMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]def gram_matrix(y):"""Input shape: b,c,h,wOutput shape: b,c,c"""(b, ch, h, w) = y.size()features = y.view(b, ch, w * h)features_t = features.transpose(1, 2)gram = features.bmm(features_t) / (ch * h * w)return gramdef get_style_data(path):"""load style image，Return： tensor shape 1*c*h*w, normalized"""style_transform = tv.transforms.Compose([tv.transforms.ToTensor(),tv.transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),])style_image = tv.datasets.folder.default_loader(path)style_tensor = style_transform(style_image)return style_tensor.unsqueeze(0)def normalize_batch(batch):"""Input: b,ch,h,w  0~255Output: b,ch,h,w  -2~2"""mean = batch.data.new(IMAGENET_MEAN).view(1, -1, 1, 1)std = batch.data.new(IMAGENET_STD).view(1, -1, 1, 1)mean = (mean.expand_as(batch.data))std = (std.expand_as(batch.data))return (batch / 255.0 - mean) / std

transformer_net.py

"""
code refer to https://github.com/abhiskk/fast-neural-style/blob/master/neural_style/transformer_net.py
"""
import torch as t
import torch.nn as nn
import numpy as npclass TransformerNet(nn.Module):#风格迁移网络def __init__(self):super(TransformerNet, self).__init__()# Down sample layersself.initial_layers = nn.Sequential(ConvLayer(3, 32, kernel_size=9, stride=1),nn.InstanceNorm2d(32, affine=True),nn.ReLU(True),ConvLayer(32, 64, kernel_size=3, stride=2),nn.InstanceNorm2d(64, affine=True),nn.ReLU(True),ConvLayer(64, 128, kernel_size=3, stride=2),nn.InstanceNorm2d(128, affine=True),nn.ReLU(True),)# Residual layersself.res_layers = nn.Sequential(ResidualBlock(128),ResidualBlock(128),ResidualBlock(128),ResidualBlock(128),ResidualBlock(128))# Upsampling Layersself.upsample_layers = nn.Sequential(UpsampleConvLayer(128, 64, kernel_size=3, stride=1, upsample=2),nn.InstanceNorm2d(64, affine=True),nn.ReLU(True),UpsampleConvLayer(64, 32, kernel_size=3, stride=1, upsample=2),nn.InstanceNorm2d(32, affine=True),nn.ReLU(True),ConvLayer(32, 3, kernel_size=9, stride=1))def forward(self, x):x = self.initial_layers(x)x = self.res_layers(x)x = self.upsample_layers(x)return xclass ConvLayer(nn.Module):"""add ReflectionPad for Conv"""def __init__(self, in_channels, out_channels, kernel_size, stride):super(ConvLayer, self).__init__()reflection_padding = int(np.floor(kernel_size / 2))self.reflection_pad = nn.ReflectionPad2d(reflection_padding)self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)def forward(self, x):out = self.reflection_pad(x)out = self.conv2d(out)return outclass UpsampleConvLayer(nn.Module):"""UpsampleConvLayerinstead of ConvTranspose2d, we do UpSample + Conv2dsee ref for why.ref: http://distill.pub/2016/deconv-checkerboard/"""def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None):super(UpsampleConvLayer, self).__init__()self.upsample = upsamplereflection_padding = int(np.floor(kernel_size / 2))self.reflection_pad = nn.ReflectionPad2d(reflection_padding)self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)def forward(self, x):x_in = xif self.upsample:x_in = t.nn.functional.interpolate(x_in, scale_factor=self.upsample)out = self.reflection_pad(x_in)out = self.conv2d(out)return outclass ResidualBlock(nn.Module):"""ResidualBlockintroduced in: https://arxiv.org/abs/1512.03385recommended architecture: http://torch.ch/blog/2016/02/04/resnets.html"""def __init__(self, channels):super(ResidualBlock, self).__init__()self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1)self.in1 = nn.InstanceNorm2d(channels, affine=True)self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1)self.in2 = nn.InstanceNorm2d(channels, affine=True)self.relu = nn.ReLU()def forward(self, x):residual = xout = self.relu(self.in1(self.conv1(x)))out = self.in2(self.conv2(out))out = out + residualreturn out

PackedVGG.py

import torch
import torch.nn as nn
from torchvision.models import vgg16
from collections import namedtupleclass Vgg16(torch.nn.Module):#Vgg16模型，已经训练好，可以分类1000种def __init__(self):super(Vgg16, self).__init__()features = list(vgg16(pretrained=True).features)[:23]# the 3rd, 8th, 15th and 22nd layer of \# self.features are: relu1_2,relu2_2,relu3_3,relu4_3self.features = nn.ModuleList(features).eval()def forward(self, x):results = []for ii, model in enumerate(self.features):x = model(x)if ii in {3, 8, 15, 22}:results.append(x)vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3'])return vgg_outputs(*results)

运行结果

H:\ProgramData\Anaconda3\python.exe D:/PycharmProjects/untitled/风格迁移/main.py
0 loss is 24709134.000000
1 loss is 24594472.000000
2 loss is 24192556.000000
3 loss is 23792200.000000Process finished with exit code -1

所想要的风格
style.png

训练6万张图片后输出：（损失从2400万降到140万）

参考文章链接
深度学习框架PyTorch入门与实践：第八章 AI艺术家：神经网络风格迁移
CNN系列学习之VGG16
训练时间10小时，完成6万张图片输入训练，模型名字中数字为读取多少张图片后生成的模型，文件夹内图片为该模型的风格，该方法一个模型只能实现一个风格，生成模型链接：https://download.csdn.net/download/qq_42017767/12821255

Pytorch笔记：风格迁移相关推荐

通过PyTorch实现风格迁移
一.环境及数据集准备 pytorch 1.1.0 torchvision 0.3.0 cuda 9.0 数据集用的是COCO2014的train2014训练集,使用ImageNet也可以需要用到在I ...
使用pytorch进行风格迁移，以艺术绘画的风格来生成一些好玩的图片
pytorch官网 https://pytorch.org/tutorials/advanced/neural_style_tutorial.html 原理风格迁移的原理基于卷积神经网络的特征提取能 ...
pytorch实现风格迁移 style transfer
本文给出简单代码实现风格迁移. 1,原理简介风格迁移和上篇文章提到的deep dream算法比较接近,都是根据某种优化指标计算梯度来反向优化输入图像的像素.所以在学完deep dream之后 ...
使用 PyTorch 进行风格迁移（Neural-Transfer）
点击上方"小白学视觉",选择加"星标"或"置顶" 重磅干货,第一时间送达 1.简介本教程主要讲解如何实现由 Leon A. Gatys,A ...
《深度学习之pytorch实战计算机视觉》第8章图像风格迁移实战（代码可跑通）
上一章<深度学习之pytorch实战计算机视觉>第7章迁移学习(代码可跑通)介绍了迁移学习.本章将完成一个有趣的应用,基于卷积神经网络实现图像风格迁移(Style Transfer).和 ...
深度学习框架PyTorch入门与实践：第八章 AI艺术家：神经网络风格迁移
本章我们将介绍一个酷炫的深度学习应用--风格迁移(Style Transfer).近年来,由深度学习引领的人工智能技术浪潮越来越广泛地应用到社会各个领域.这其中,手机应用Prisma,尝试为用户的照片 ...
Neural Style Transfer 风格迁移经典论文讲解与 PyTorch 实现
今天花半小时看懂了"Image Style Transfer Using Convolutional Neural Networks Leon"这篇论文,又花半小时看懂了其 PyT ...
李沐d2l《动手学深度学习》第二版——风格迁移源码详解
本文是对李沐Dive to DL<动手学深度学习>第二版13.12节风格迁移的源码详解,整体由Jupyter+VSCode完成,几乎所有重要代码均给出了注释,一看就懂.需要的同学可以在文末 ...
Pytorch入门+实战系列七：图片风格迁移和GAN
Pytorch官方文档:https://pytorch.org/docs/stable/torch.html? 1. 写在前面今天开始,兼顾Pytorch学习, 如果刚刚接触深度学习并且想快速搭建神 ...

Pytorch笔记：风格迁移

Pytorch笔记：风格迁移相关推荐

最新文章

热门文章