通過PyTorch實現風格遷移

一.環境及數據集準備

pytorch 1.1.0

torchvision 0.3.0

cuda 9.0

數據集用的是COCO2014的train2014訓練集,使用ImageNet也可以

需要用到在ImageNet上預訓練好的VGG16,模型文件爲vgg16-397923af.pth

二.原理

​ 首先需要明白一點,深度學習之所以被稱爲“深度”,就在於它採用了深層的網絡結構,網絡的不同層學到的是圖像不同層面的特徵信息

​ 研究表明,幾乎所有神經網絡的第一層學習到的都是關於線條和顏色的信息(直觀理解就是像素組成色彩,點組成線)。再往上,神經網絡開始關注一些複雜的特徵(比如拐角或特殊形狀等),這些特徵可以看成是低層次的特徵組合。隨着深度的加深,神經網絡關注的信息逐漸抽象(例如有的卷積核關注嘴巴,有的關注眼睛等,以及關注對象之間的空間關係等)。

​ 在進行風格遷移時,並不要求生成圖片的像素和輸入圖片中每一個像素都一樣,而追求的是生成圖片和輸入圖片具有相同的特徵(比如原圖是個人,我們希望風格遷移之後的圖片依據能明顯看出來是個人)。但這是不是就類似神經網絡中最後的層的輸出(差不多也就是最後的卷積核關注的特徵)?但是這最後一層的特徵的抽象程度太高,我們不可能只希望最後只從生成圖片中看出來是個人,我們很是希望生成圖片中能保存輸入圖片的部分細節信息(五官、動作等),這些信息相對來說抽象程度沒那麼高。因此我們使用中間某些層的特徵作爲目標,希望輸入圖片和風格遷移的生成圖片在這些層的特徵儘可能的相似,即將圖片在深度模型中的中間某些層的輸出作爲圖像的的知覺特徵

​ 注:這就是下面要使用VGG-16中不同激活層(ReLU)輸出的原因所在

一般使用Gram矩陣來表示圖像的風格特徵。對於每一張圖片,卷積層的輸出形狀爲C×H×W,C是卷積核的通道數,一般稱爲有C個卷積核,每個卷積核學習圖像的不同特徵。每一個卷積核輸出的H×W代表這張圖像的一個feature map(特徵圖),最後所有的卷積和輸出組合成feature maps。通過計算每個feature map之間的相似性,可以得到圖像的風格特徵。

​ 對於一個C×H×W的feature maps F ,Gram Matrix的形狀爲C×C,其第i、j個元素Gi,j的計算方式定義如下:
Gi,j=kFikFjk G_{i,j} = \sum_kF_{ik}F_{jk}
​ 其中Fik代表第i個feature map的第k個像素點。

​ 關於Gram Matrix,要注意:

  • Gram Matrix的計算採用累積的形式,拋棄了空間信息。一張圖片的像素隨機打亂之後計算得到的Gram Matrix和原圖的Gram Matrix一樣。所以可以認爲Gram Matrix拋棄了元素之間的空間信息。
  • Gram Matrix的結果與feature maps F的尺度無關,只與通道數有關。無論H、W的大小如何,最後Gram Matrix的形狀都是C×C。
  • 對於一個C×H×W的feature maps,可以通過調整形狀和矩陣乘法快速計算它的Gram Matrix,即先將F調整爲C×(HW)的二維矩陣,然後再計算F×F^T,結果就是Gram Matrix。(這是因爲矩陣C×(HW)的轉置矩陣爲(HW)×C,所以矩陣乘法C×(HW) * (HW)×C的結果就是C×C。)

​ 實踐證明Gram Matrix的特點:注重風格紋理等特徵,忽略空間信息。圖像的空間信息在計算Gram Matrix時都被捨棄,但是紋理、色彩等風格信息被保存下來。Gram Matrix表徵圖像的風格特徵在風格遷移、紋理合成等任務中的表現十分出衆。

​ 風格遷移的圖片想要做到逼真。第一是要生成的圖片在內容、細節上儘可能地與輸入的內容圖片相似;第二要生成的圖片在風格上儘可能地與風格圖片相似。因此,下面會用損失content loss和style loss來衡量這兩個指標。

​ 在下面我們會實現一種快速風格遷移算法——Fast Style Transfer,其更多被稱爲Fast Neural Style

​ Fast Neural Style專門設計了一個網絡用來進行風格遷移,輸入原圖片,網絡將自動生成目標圖片。這個網絡會將輸入的風格圖片訓練成一個相對應的風格網絡,接下來輸入我們想進行轉換的圖片,很快就能完成一次風格遷移。

​ 在Fast Neural Style的網絡結構中,x是輸入圖像。在風格遷移任務中yc=xys是風格圖片,Image Transform Net fw是我們設計的風格遷移網絡,針對輸入的圖像x,能夠返回一張新的圖像y**。**y在圖像內容上與yc相似,但在風格上與ys相似。損失網絡(Loss Network)不用訓練,只用來計算知覺特徵和風格特徵。損失網絡,採用在ImageNet上預訓練好的VGG-16。

圖片來自於參考中的第一個

​ VGG-16網絡,從上到下有5個卷積塊,兩個卷積塊間通過MaxPooling層區分。每個卷積塊有2~3個卷積層,每一個卷積層後面都跟着一個ReLU激活層。

​ Fast Neural Style的訓練步驟:

​ (1)輸入一張圖片xfw(風格遷移網絡) 中得到 y^(生成的圖片)。

​ (2)將 y^yc(就是x) 輸入到Loss Network(VGG-16)中,計算它在relu3_3(表示第3個卷積塊的第3個卷積層的,激活層)的輸出,並計算它們之間的均方誤差作爲content loss(生成的圖片在內容、細節上與輸入的內容圖片的相似程度)。

​ (3)將 y^ys(風格圖片)輸入到Loss Network中,計算它在 relu1_2relu2_2relu3_3relu4_4 的輸出,再計算它們的Gram Matrix 的均方誤差作爲 style loss(生成的圖片在風格上與風格圖片的相似程度)。

​ (4)兩個損失相加,再反向傳播。更新 fw 的參數,固定Loss Network不動(即不更新損失網絡的參數)。

​ (5)重複(1)到(4),訓練 fw

三.使用PyTorch實現風格遷移

1.項目的文件組織目錄

  • checkpoints/:用來保存模型。
  • data/:用來保存數據,可以通過直接或軟連接的方式將數據保存在data文件夾下,這樣方便之後通過torchvision的ImageFolder方法直接加載文件夾中的圖片。
  • main.py:主函數,包括訓練和測試。
  • PackedVGG.py:預訓練好的VGG-16,爲了提取中間層的輸出,做了一些簡化。
  • transformer_net.py:風格遷移網絡。輸入一張圖片,輸出一張圖片。
  • utils.py:工具包,主要是計算Gram Matrix、加載風格圖片和batch數據標準化等。

2.預訓練的VGG-16

VGG16的網絡結構圖如下。

VGG-16

PackedVGG.py

#coding:utf-8
 
import torch
import torch.nn as nn
from torchvision.models import vgg16
from collections import namedtuple

class Vgg16(nn.Module):
	def __init__(self):
		super(Vgg16, self).__init__()
		features = list(vgg16(pretrained=True).features[:23]) # 把預訓練的vgg16的前23層放在一個列表裏,vgg16總共31層(除了全連接層)
		# features的第3,8,15,22層分佈是:relu1_2,relu2_2,relu3_3,relu4_3
		self.features = nn.ModuleList(features).eval() # 不更新參數

	def forward(self, x):
		results = [] # 結果列表
		for ii,model in enumerate(self.features): # 獲取3,8,15,22層的relu的輸出
			x = model(x)
			if ii in {3, 8, 15, 22}:
				results.append(x)

		# namedtuple創建一個具名元組,參數有,typename代表元組名稱,field_names代表元組中元素的名稱,rename當元素名稱中有Python的關鍵字,則必須設置爲True
		vgg_outputs = namedtuple("VggOutputs",['relu1_2','relu2_2','relu3_3','relu4_3'])
		# * + results表示把列表中的數據全部提出來,等於去掉最外層的[]
		return vgg_outputs(*results) # 返回一個具名元組,元組的元素是3,8,15,22relu層的輸出

上邊我們在網絡前向傳播的過程中獲得了中間層的輸出,並保存下來,其他的層不再需要,節省了空間。

​ 在torchvision中,VGG的實現由兩個nn.Sequential對象組成。第一個是features,包括卷積、激活和池化等層,用來提取圖像特徵。另一個是classifier,包含全連接層,用來分類。上邊通過vgg.features直接獲得對應的nn.Sequential對象。這樣在前向傳播時,當計算完指定層的輸出後,就將結果保存於一個list中,然後再使用namedtuple進行名稱綁定,這樣可以通過output.relu1_2訪問第一個元素,方便又直觀。當然也可以利用layer.register_forward_hook的方式來獲取相應層的輸出,但是在本例中相對比較麻煩。

3.實現風格遷移網絡

​ 風格遷移網絡,參考了PyTorch的官方示例,其網絡結構如下圖所示。

圖片來自於參考中的第一個

​ 圖中(a)是網絡的總體結構,左邊(d)是一個殘差單元的結構圖,右邊(b)和(c)分別是下采樣和上採樣單元的結構圖。網絡結構總結有以下幾個特點。

  • 先下采樣,然後上採樣,這種做法使計算量變小。
  • 使用殘差結構使網絡變深。
  • 邊緣補齊的方式不再是傳統的補0,而是採用一種被稱爲Reflection Pad的補齊策略:上下左右反射邊緣的像素進行補齊。
  • 上採樣不再使用傳統的ConvTransposed2d,而是先用Upsample,然後用Conv2d,這中做法能避免Checkerboard Artifacts現象棋盤效應)。
  • Batch Normalization全部改成Instance Normalization
  • 網絡中沒有全連接層,線性操作是卷積,因此對輸入和輸出圖片的尺寸沒有要求,這裏我們輸入和輸出圖片的尺寸都是3×256×256(其他尺寸的一樣可以,可以自定義做修改)。

transformer_net.py

#coding:utf-8
 
import torch as t
from torch import nn
import numpy as np


class TransformerNet(nn.Module):
	"""
	InstanceNorm:一個channel內做歸一化,算H*W的均值,用在風格化遷移;
	因爲在圖像風格化中,生成結果主要依賴於某個圖像實例,所以對整個batch歸一化不適合圖像風格化中,因而對HW做歸一化。可以加速模型收斂,並且保持每個圖像實例之間的獨立。
	affine: 布爾值,當設爲true,給該層添加可學習的仿射變換參數。
	仿射變換即進行平移、旋轉、轉置、縮放等操作。

	"""
	def __init__(self):
		super(TransformerNet, self).__init__()

		# 下卷積層
		self.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), # stride爲2,使得特徵圖尺寸縮小1/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 layers(殘差層)
        self.res_layers = nn.Sequential(
			ResidualBlock(128),
			ResidualBlock(128),
			ResidualBlock(128),
			ResidualBlock(128),
			ResidualBlock(128)
		)

		# Upsampling Layers(上卷積層)
		self.upsample_layers = nn.Sequential(
			UpsampleConvLayer(128, 64, kernel_size=3, stride=1, upsample=2), # upsample取2是因爲stride取1,爲了使特徵圖的尺寸擴大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):
		out = self.initial_layers(x) # 下采樣的卷積層
		out = self.res_layers(out) # 深度殘差層
		out = self.upsample_layers(out) # 上採樣的卷積層
		return out

    
# 卷積模塊
class 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)) # np.floor 返回不大於輸入參數的最大整數。

		# 定義反射填充層  H(out) = H(in) + paddingTop + paddingBottom 
        # 高加上下,寬加左右,平常我們的tensor是(B,C,H,W) 
		self.reflection_pad = nn.ReflectionPad2d(reflection_padding) # ReflectionPad2d的參數有四個,比如傳一個元組(3,3,5,5) last right top bottom,只傳一個數表示上下左右都是同一個數

		# 定義卷積層
		self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)

	def forward(self, x):
		# 在經過卷積層前先給x加個反射填充層,例如例如,(64,3,220,220) 過(3,3,5,5)得到(64,3,230,226),似乎我們一般都是上下左右加同一個數,做到鏡像
		out = self.reflection_pad(x) # 這裏上下左右加同一個數,使得特徵圖的H和W增大了
		out = self.conv2d(out)
		return out
    
    
# 上採樣模塊
class UpsampleConvLayer(nn.Module):
	"""
	這裏也使用了邊界反射填充
	先上採樣,然後做一個卷積(Conv2d),而不是採用ConvTranspose2d,這種效果更好,爲了避免棋盤效應
	比如,這裏的kernel_size取3,stride取1,upsample取2,輸入x的尺寸爲2×2
	那麼,經過nn.functional.interpolate()後特徵圖變成4×4
	經過reflection_pad()後特徵圖變成6×6,再通過conv2d()後特徵圖恢復到4×4
	
	"""
	def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None):
		super(UpsampleConvLayer, self).__init__()

		self.upsample = upsample # 上採樣的標誌,用來表示特徵圖擴大的倍數

		# 下面同上面的ConvLayer
		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):
		if self.upsample: # 判斷是否上採樣
			# interpolate實現插值和上採樣,插值算法取決於參數mode的設置,默認爲最近鄰nearest,scale_factor指定輸出爲輸入的多少倍數
			x = t.nn.functional.interpolate(x, scale_factor=self.upsample)
			
		out = self.reflection_pad(x)
		out = self.conv2d(out)
		return out

    
# 殘差層
class ResidualBlock(nn.Module):
    """
    這裏kernel_size取3,使得reflection_padding反射填充的層數取1,
    通過reflection_pad層就會給原特徵圖H和C各增加2,
    再經過一個(k=3,p=0,s=1)的卷積後特徵圖的H和C又各減去2,
    因此維持了經過self.conv1層的特徵圖的尺寸不變。
    由此,可以明白,這裏的殘差網絡加深了網絡,且沒有改變原本特徵圖的尺寸。
    """
	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):
		identity = x # 先將本體保存下來
		out = self.conv1(x)
		out = self.in1(out)
		out = self.relu(out)
		out = self.conv2(out)
		out = self.in2(out)
		out = out + identity # 將原本的x和最後的out相加在一次,實現殘差
		return out # 這裏或許可以再通過一個激活層,這樣self.relu(out),但我沒有這樣實驗過

​ **nn.functional.interpolate()**在插值算法上可以使用的上採樣算法,有’nearest’,‘linear’,‘bilinear’,‘bicubic’,‘trilinear’和’area’,默認採用’nearest’。

​ 可以這樣直觀的觀察interpolate()是如何使用的

input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
input
tensor([[[[1., 2.],
          [3., 4.]]]])

x = F.interpolate(input, scale_factor=2, mode='nearest')
x
tensor([[[[1., 1., 2., 2.],
          [1., 1., 2., 2.],
          [3., 3., 4., 4.],
          [3., 3., 4., 4.]]]])

x = F.interpolate(input, scale_factor=2, mode='bilinear', align_corners=True)
x
tensor([[[[1.0000, 1.3333, 1.6667, 2.0000],
          [1.6667, 2.0000, 2.3333, 2.6667],
          [2.3333, 2.6667, 3.0000, 3.3333],
          [3.0000, 3.3333, 3.6667, 4.0000]]]])

想了解PyTorch實現上採樣的方法和原理,可以參考這個博客

4.工具包

工具包裏實現了計算Gram Matrix加載風格圖片batch數據標準化

我們還可以在其中實現用於可視化的工具。

utils.py

#coding:utf-8
 
# 工具類
from itertools import chain
import torch as t
import torchvision as tv
import numpy as np

# 這裏的均值和標準差不是0.5和0.5,而是使用他人專門計算的ImageNet上所有圖片的均值和標準差
# 這更符合真實世界圖片的分佈,效果也比都爲0.5好
IMAGENET_MEAN = [0.485,0.456,0.406] # 均值
IMAGENET_STD = [0.229,0.224,0.225] # 方差

def gram_matrix(y):
	"""
	使用Gram矩陣來表示圖像的風格特徵
	輸入B,C,H,W
	輸出B,C,C

	"""
	(b,ch,h,w) = y.size() # 比如1,8,2,2
	features = y.view(b,ch,w*h)  # 得到1,8,4
	features_t = features.transpose(1,2) # 調換第二維和第三維的順序,即矩陣的轉置,得到1,4,8
	gram = features.bmm(features_t) / (ch*h*w) # bmm()用來做矩陣乘法,及未轉置的矩陣乘以轉置後的矩陣,得到的就是1,8,8了
    # 由於要對batch中的每一個樣本都計算Gram Matrix,因此使用bmm()來計算矩陣乘法,而不是mm()
	return gram


def get_style_data(path):
    """
    加載風格圖片,
    輸入: path, 文件路徑
    返回: tensor 形狀爲1*c*h*w, 分佈大約在-2~2
    """
    # 數據預處理
    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) # 圖片的tensor形狀爲 c*h*w
    return style_tensor.unsqueeze(0) # 增加一維,變成1*c*h*w


def normalize_batch(batch):
    """
    batch數據標準化
    輸入: b,c,h,w  將分佈在0~255的圖片進行標準化
    輸出: b,c,h,w  大約-2~2
    """
    mean = batch.data.new(IMAGENET_MEAN).view(1, -1, 1, 1) # .new()等於創建type和device與batch.data一樣的tensor,不帶參的new()無內容
    std = batch.data.new(IMAGENET_STD).view(1, -1, 1, 1)
    mean = (mean.expand_as(batch.data)) # 表示讓mean變成和batch.data一樣形狀的tensor
    std = (std.expand_as(batch.data)) 
    return (batch / 255.0 - mean) / std # 返回標準化處理後的batch
上邊new()中給參數,是在添加噪聲,並且將數據維度變成(1,-1,1,1),比如64×64就變成了1×4096×1×1,其他維度變成1,那麼-1位置就要增加其他維度的值

​ 比如 (3,1).expand_as(3,4) 形狀就從(3,1)變成了(3,4),不夠的數據,就重複擴充,比如**[[1],[2],[3]] -> [[1,1,1,1],[2,2,2,2],[3,3,3,3]]**。這裏是爲了讓單獨的mean和std擴充成batch形狀,方便之後batch中每張圖片的標準化,這樣矩陣的計算就省去了單獨一張圖片一張圖片的標準化了。

5.訓練和測試

main.py

#coding:utf-8

# 風格遷移的主程序
import torch as t
import torchvision as tv
from torch.utils import data
from torch.nn import functional as F
import tqdm
import time
import os

mean = [0.485,0.456,0.406] # 平均
std = [0.229,0.224,0.225] # 方差

# 定義各項參數類
class Config:
	image_size = 256 # 圖片大小
	batch_size = 8 # 批處理數
	data_root = 'data/coco2014/train2014/'
	num_workers = 4 # 多線程加載數據
	use_gpu = True # 默認使用GPU

	style_path = 'style.jpeg' # 風格圖片存放的路徑

	lr = 1e-3 # 學習率
	plot_every = 100 # 每10個batch可視化一次

	epoches = 2 # 訓練輪數

	content_weight = 1e5 # content_loss的權重
	style_weight = 1e10 # style_loss的權重

	model_path = None # 預訓練模型的路徑
	save_path = 'imgs'
	content_path = 'content.jpg' # 需要進行風格遷移的圖片
	result_path = 'output.png' # 風格遷移後結果的保存路徑
    
    
# 定義訓練過程
def train(**kwargs):
	opt = Config() # 加載參數

	# 加載數據和數據預處理
	transforms = tv.transforms.Compose([
		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,transforms)
	dataloader = data.DataLoader(dataset,opt.batch_size,num_workers=opt.num_workers,drop_last=True)
	print("Dataset:",len(dataset),"Dataloader:",len(dataloader)) 
    # Dataset: 82783 Dataloader: 10347
    
	# 轉換網絡
	transformer = TransformerNet().cuda()
	if opt.model_path:
		print("轉換模型參加加載成功!")
		transformer.load_state_dict(t.load(opt.model_path,map_location=lambda _s,_:_s))

	# 損失網絡Vgg16
	vgg = Vgg16().cuda().eval() 
	for param in vgg.parameters():
		param.requires_grad = False # 不反向傳播更新vgg的參數

	# 優化器
	optimizer = t.optim.Adam(transformer.parameters(),opt.lr)


	# 獲取風格圖片的數據
	style = get_style_data(opt.style_path)
	style = style.cuda()
    
    # 風格圖片的gram矩陣
	with t.no_grad():
		features_style = vgg(style) # vgg()只返回了 3,8,15,22的relu層的輸出
		gram_style = [gram_matrix(y) for y in features_style] # 分別計算這些層的Gram Matrix
    
	now = time.time()
	for epoch in range(opt.epoches):
		c_loss = 0
		s_loss = 0
		for ii,(x,_) in tqdm.tqdm(enumerate(dataloader)):
			# 訓練
			optimizer.zero_grad() # 梯度清零
			x = x.cuda()
			y = transformer(x) # 得到風格遷移後的圖片
            # 對x和y(即yc和y^)進行batch標準化處理
			y = normalize_batch(y)
			x = normalize_batch(x)
            # 獲得x和y在vgg上那四層的輸出
			features_y = vgg(y)
			features_x = vgg(x)

			# content loss 內容loss,只用到了relu3_3(在原理中有介紹)
			content_loss = opt.content_weight * F.mse_loss(features_y.relu3_3,features_x.relu3_3)

			# style loss 風格loss,這裏四層relu都用到了
			style_loss = 0
			for ft_y,gm_s in zip(features_y,gram_style):
				gram_y = gram_matrix(ft_y)
				style_loss += F.mse_loss(gram_y,gm_s.expand_as(gram_y))
			style_loss *= opt.style_weight

			total_loss = content_loss + style_loss
			total_loss.backward() # 反向傳播
			optimizer.step() # 更新參數

			# 損失累加
			c_loss += content_loss.item()
			s_loss += style_loss.item()
            
            if ii%opt.plot_every == 0:
				# 每隔一定的batch保存風格遷移後的圖片,也可以順便把原圖片也保存下來,方便對比
				tv.utils.save_image((y.data.cpu()[0]*0.225+0.45).clamp(min=0,max=1),'%s/%s_output_%s.png'%(opt.save_path,epoch,ii),normalize=True,range=(-1,1))
				# tv.utils.save_image((x.data.cpu()[0]*0.225+0.45).clamp(min=0,max=1),'%s/%s_input_%s.png'%(opt.save_path,epoch,ii),normalize=True,range=(-1,1))
			if ii%1000 == 0: # 每隔一定batch打印一次訓練信息
				print("Epoch:{},C_Loss:{:.6f},S_Loss:{:.6f},Time:{:.4f}s".format(epoch,c_loss/8000,s_loss/8000,time.time()-now))
                
        # 保存模型
		t.save(transformer.state_dict(),'checkpoints/%s_style.pth'%epoch)

# 開始訓練 
train()


# 測試
def stylize(**kwargs):
	with t.no_grad():
		opt = Config()

		# 加載要進行風格遷移的圖片,並進行預處理
		content_image = tv.datasets.folder.default_loader(opt.content_path)
		content_transform = tv.transforms.Compose([
			tv.transforms.ToTensor(),
			tv.transforms.Lambda(lambda x:x.mul(255)) # 使用transforms.Lambda封裝其爲transforms策略,mul()也是矩陣的一種點乘操作,要求操作的兩個矩陣維度必須一樣
		])
		content_image = content_transform(content_image)
		content_image = content_image.unsqueeze(0).cuda().detach() # upsqueeze(0)增加了一個維度,使c×h×w變成了1×c×h×w

		# 模型
		style_model = TransformerNet().eval() # map_location表示將GPU保存的模型加載到CPU上
		style_model.load_state_dict(t.load(opt.model_path,map_location=lambda _s,_:_s))
		style_model = style_model.cuda() # 或這樣map_location=lambda storage, loc: storage

		# 風格遷移與保存
		output = style_model(content_image)
		output_data = output.cpu().data[0]
		tv.utils.save_image(((output_data/255)).clamp(min=0,max=1),opt.result_path)

stylize()

爲上面代碼說明幾點:

  • 圖片的每個像素的取值範圍爲0~255。
  • 調用torchvision的transforms.ToTensor()操作,像素會被轉換到0~1。
  • 標準化(減均值後除以標準差),這裏使用的均值和標準差爲[0.485,0.456,0.406]和[0.229,0.224,0.225]。可以估算這時圖片的分佈範圍大概在 (0-0.4845)/0.229≈**-2.1和 (1-0.406)/0.225≈2.7**之間。儘管這是它的分佈在-2.1~2.7,但是它的均值接近0,標準差接近1,採用ImageNet圖片的均值和標準差作爲標準化參數的目的是圖像各個像素的分佈接近標準分佈。
  • VGG-16網絡的輸入圖像數值大小爲使用ImageNet均值和標準差進行標準化之後的圖片數據,即-2.1~2.7。
  • TransformerNet網絡的輸入圖片的像素值是0255,輸出的像素值也希望是0255,但是由於輸出沒有做特殊處理,所以可能出現小於0和大於255的像素。

​ 另外除了直接調用train()和stylize()還可以通過終端的方式改變Config()類中的配置的值的方式來靈活運行訓練和測試。

if __name__ == '__main__':
    import fire
	
	fire.Fire()
    
# 這樣寫需要先註釋掉上面的train()和stylize()
# 然後通過終端來運行訓練和測試過程,例如
python main.py train --style_path=XXX.jpeg

python main.py stylize --model_path='checkpoints/XXX.pth' --content_path=XXX.jpg --result_path=XXX.png --use_gpu=False # 這裏字符串的引號是可以省略的

6.結果分析

風格圖片1

風格圖片1

原圖1
原圖1

原圖2
原圖2

遷移後的圖片
風格遷移圖片1風格遷移圖片2
可以看到效果很是不錯,風格圖片的風格特徵被很好的學習到,且原圖的知覺特徵也很明顯。

下面換一個風格圖片,梵高的風格

風格圖片2

遷移後的圖片:
風格遷移圖片3風格遷移圖片4

也還不錯,把風格圖片的顏色搭配學到了,不過藝術風格還差點意思。大家可以嘗試訓練其他風格的圖片。
另外需要提出的一點,上面訓練2個epoch花費了大約1.6個小時,是放在Kaggle上訓練的,其用的GPU應該是Nvida Tesla P100。這不用我吹了,可以想想有多香把。

四.總結

​ 風格遷移非常的有趣,可以把我們的圖片轉換成我們喜歡的風格,且快速完成轉換的過程。

​ 我們上面實現的是Fast Neural Style,在其之前還有風格遷移開山之作Neural Style,但是Neural Style在進行風格遷移時需要花費幾十分鐘甚至幾個小時的訓練,但是其生成圖片的效果會比Fast Neural Style好上一點。不過,速度是個大硬傷,我們總喜歡很快就能看到滿意的結果。

​ 除了Neural Style和Fast Neural Style,還有別的有趣又吸引人的風格遷移項目。

​ Adobe的圖片風格深度遷移(Deep Photo Style Transfer),還有CycleGAN,它們在風格遷移上表現的尤爲出色,比Fast Neural Style的效果要好上許多。其中,CycleGAN的網絡結構和Fast Neural Style的transformer類似,但它採用的是GAN的訓練方式,能夠實現風格的雙向轉換(即原圖片轉換成風格圖片,風格圖片轉換成原圖片的風格)。這個我會在接下來抽空去接觸並實踐一下,有趣的東西還是要多取嘗試的。

五.參考

深度學習框架PyTorch入門與實踐(https://github.com/chenyuntc/pytorch-book/tree/master/chapter8-%E9%A3%8E%E6%A0%BC%E8%BF%81%E7%A7%BB(Neural%20Style))

Deconvolution and Checkerboard Artifacts(https://www.jianshu.com/p/36ff39344de5)

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章