樣式遷移

將樣式圖片中的風(fēng)格遷移到任意圖片上，比使用濾鏡更加靈活高效。

基于CNN的樣式遷移：

以下圖為例，圖中展示的是一個(gè)三層CNN網(wǎng)絡(luò)，（為了便于理解畫成三個(gè)CNN）CNN實(shí)現(xiàn)樣式遷移的方式是，讓輸入圖片在某些CNN層與內(nèi)容圖片進(jìn)行匹配，在另外的一些CNN層與樣式圖片進(jìn)行匹配。從而使得最終輸出的圖片既有內(nèi)容圖片的內(nèi)容，也具備樣式圖片的特征。

該算法的特點(diǎn)是，訓(xùn)練對(duì)象并不是網(wǎng)絡(luò)中的參數(shù)，而是對(duì)輸入的圖片進(jìn)行訓(xùn)練，使得該圖片越來(lái)越接近我們想要得到的效果。

首先，我們讀取內(nèi)容和風(fēng)格圖像。 從打印出的圖像坐標(biāo)軸可以看出，它們的尺寸并不一樣。

%matplotlib inline
import torch
import torchvision
from torch import nn
from d2l import torch as d2l

d2l.set_figsize()
content_img = d2l.Image.open('../img/rainier.jpg')
d2l.plt.imshow(content_img);

style_img = d2l.Image.open('../img/autumn-oak.jpg')
d2l.plt.imshow(style_img);

下面，定義圖像的預(yù)處理函數(shù)和后處理函數(shù)。 預(yù)處理函數(shù)preprocess對(duì)輸入圖像在RGB三個(gè)通道分別做標(biāo)準(zhǔn)化，并將結(jié)果變換成卷積神經(jīng)網(wǎng)絡(luò)接受的輸入格式。 后處理函數(shù)postprocess則將輸出圖像中的像素值還原回標(biāo)準(zhǔn)化之前的值。 由于圖像打印函數(shù)要求每個(gè)像素的浮點(diǎn)數(shù)值在0～1之間，我們對(duì)小于0和大于1的值分別取0和1。

rgb_mean = torch.tensor([0.485, 0.456, 0.406])
rgb_std = torch.tensor([0.229, 0.224, 0.225])

def preprocess(img, image_shape):
    transforms = torchvision.transforms.Compose([
        torchvision.transforms.Resize(image_shape),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=rgb_mean, std=rgb_std)])
    return transforms(img).unsqueeze(0)

def postprocess(img):
    img = img[0].to(rgb_std.device)
    img = torch.clamp(img.permute(1, 2, 0) * rgb_std + rgb_mean, 0, 1)
    return torchvision.transforms.ToPILImage()(img.permute(2, 0, 1))

使用基于ImageNet數(shù)據(jù)集預(yù)訓(xùn)練的VGG-19模型來(lái)抽取圖像特征?(Gatys?et al., 2016)。

VGG-19對(duì)于抽取特征任務(wù)來(lái)說(shuō)是比較合適的。

pretrained_net = torchvision.models.vgg19(pretrained=True)

指定樣式層和內(nèi)容層（層數(shù)越小越靠近輸入，越大越靠近輸出）

樣式層既需要對(duì)圖片的細(xì)節(jié)（例如筆觸）進(jìn)行提取，又需要對(duì)圖片的整體風(fēng)格進(jìn)行把握（例如圖畫色彩基調(diào)）所以層數(shù)取得相對(duì)平均。

而內(nèi)容層則取靠上的位置，允許為了樣式匹配而稍微修改圖片內(nèi)容。

style_layers, content_layers = [0, 5, 10, 19, 28], [25]

使用VGG層抽取特征時(shí)，我們只需要用到從輸入層到最靠近輸出層的內(nèi)容層或風(fēng)格層之間的所有層。 下面構(gòu)建一個(gè)新的網(wǎng)絡(luò)net，它只保留需要用到的VGG的所有層。

net = nn.Sequential(*[pretrained_net.features[i] for i in
                      range(max(content_layers + style_layers) + 1)])

給定輸入X，如果我們簡(jiǎn)單地調(diào)用前向傳播net(X)，只能獲得最后一層的輸出。 由于我們還需要中間層的輸出，因此這里我們逐層計(jì)算，并保留內(nèi)容層和風(fēng)格層的輸出。

def extract_features(X, content_layers, style_layers):
    contents = []
    styles = []
    for i in range(len(net)):
        X = net[i](X)
        if i in style_layers:
            styles.append(X)
        if i in content_layers:
            contents.append(X)
    return contents, styles

下面定義兩個(gè)函數(shù)：get_contents函數(shù)對(duì)內(nèi)容圖像抽取內(nèi)容特征；?get_styles函數(shù)對(duì)風(fēng)格圖像抽取風(fēng)格特征。 因?yàn)樵谟?xùn)練時(shí)無(wú)須改變預(yù)訓(xùn)練的VGG的模型參數(shù)，所以我們可以在訓(xùn)練開始之前就提取出內(nèi)容特征和風(fēng)格特征。 由于合成圖像是風(fēng)格遷移所需迭代的模型參數(shù)，我們只能在訓(xùn)練過(guò)程中通過(guò)調(diào)用extract_features函數(shù)來(lái)抽取合成圖像的內(nèi)容特征和風(fēng)格特征。

def get_contents(image_shape, device):
    content_X = preprocess(content_img, image_shape).to(device)
    contents_Y, _ = extract_features(content_X, content_layers, style_layers)
    return content_X, contents_Y

def get_styles(image_shape, device):
    style_X = preprocess(style_img, image_shape).to(device)
    _, styles_Y = extract_features(style_X, content_layers, style_layers)
    return style_X, styles_Y

定義損失函數(shù)

內(nèi)容損失

與線性回歸中的損失函數(shù)類似，內(nèi)容損失通過(guò)平方誤差函數(shù)衡量合成圖像與內(nèi)容圖像在內(nèi)容特征上的差異。 平方誤差函數(shù)的兩個(gè)輸入均為extract_features函數(shù)計(jì)算所得到的內(nèi)容層的輸出。

def content_loss(Y_hat, Y):
    # 我們從動(dòng)態(tài)計(jì)算梯度的樹中分離目標(biāo)：
    # 這是一個(gè)規(guī)定的值，而不是一個(gè)變量。
    return torch.square(Y_hat - Y.detach()).mean()

在定義樣式損失函數(shù)前，先要討論一下如何量化“樣式”這個(gè)概念。

可用的一種方法是通過(guò)統(tǒng)計(jì)圖片中RGB的直方圖來(lái)表示。通過(guò)統(tǒng)計(jì)一個(gè)通道內(nèi)的RGB分布和不同通道間的RGB分布來(lái)定義樣式。

樣式相似的兩張圖片并不是每個(gè)像素的值都一樣，而是通道之間與通道之內(nèi)的值的統(tǒng)計(jì)分布能夠相匹配。像素決定內(nèi)容，整體分布決定樣式。

此處使用格拉姆矩陣來(lái)實(shí)現(xiàn)圖像風(fēng)格的輸出

def gram(X):
    num_channels, n = X.shape[1], X.numel() // X.shape[1]
    X = X.reshape((num_channels, n))
    return torch.matmul(X, X.T) / (num_channels * n)

自然地，風(fēng)格損失的平方誤差函數(shù)的兩個(gè)格拉姆矩陣輸入分別基于合成圖像與風(fēng)格圖像的風(fēng)格層輸出。這里假設(shè)基于風(fēng)格圖像的格拉姆矩陣gram_Y已經(jīng)預(yù)先計(jì)算好了。

def style_loss(Y_hat, gram_Y):
    return torch.square(gram(Y_hat) - gram_Y.detach()).mean()

def tv_loss(Y_hat):
    return 0.5 * (torch.abs(Y_hat[:, :, 1:, :] - Y_hat[:, :, :-1, :]).mean() +
                  torch.abs(Y_hat[:, :, :, 1:] - Y_hat[:, :, :, :-1]).mean())

風(fēng)格轉(zhuǎn)移的損失函數(shù)是內(nèi)容損失、風(fēng)格損失和總變化損失的加權(quán)和。 通過(guò)調(diào)節(jié)這些權(quán)重超參數(shù)，我們可以權(quán)衡合成圖像在保留內(nèi)容、遷移風(fēng)格以及去噪三方面的相對(duì)重要性。

content_weight, style_weight, tv_weight = 1, 1e3, 10

def compute_loss(X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram):
    # 分別計(jì)算內(nèi)容損失、風(fēng)格損失和全變分損失
    contents_l = [content_loss(Y_hat, Y) * content_weight for Y_hat, Y in zip(
        contents_Y_hat, contents_Y)]
    styles_l = [style_loss(Y_hat, Y) * style_weight for Y_hat, Y in zip(
        styles_Y_hat, styles_Y_gram)]
    tv_l = tv_loss(X) * tv_weight
    # 對(duì)所有損失求和
    l = sum(10 * styles_l + contents_l + [tv_l])
    return contents_l, styles_l, tv_l, l

在風(fēng)格遷移中，合成的圖像是訓(xùn)練期間唯一需要更新的變量。因此，我們可以定義一個(gè)簡(jiǎn)單的模型SynthesizedImage，并將合成的圖像視為模型參數(shù)。模型的前向傳播只需返回模型參數(shù)即可。

class SynthesizedImage(nn.Module):
    def __init__(self, img_shape, **kwargs):
        super(SynthesizedImage, self).__init__(**kwargs)
        self.weight = nn.Parameter(torch.rand(*img_shape))

    def forward(self):
        return self.weight

下面，我們定義get_inits函數(shù)。該函數(shù)創(chuàng)建了合成圖像的模型實(shí)例，并將其初始化為圖像X。風(fēng)格圖像在各個(gè)風(fēng)格層的格拉姆矩陣styles_Y_gram將在訓(xùn)練前預(yù)先計(jì)算好。

def get_inits(X, device, lr, styles_Y):
    gen_img = SynthesizedImage(X.shape).to(device)
    gen_img.weight.data.copy_(X.data)
    trainer = torch.optim.Adam(gen_img.parameters(), lr=lr)
    styles_Y_gram = [gram(Y) for Y in styles_Y]
    return gen_img(), styles_Y_gram, trainer

訓(xùn)練模型

在訓(xùn)練模型進(jìn)行風(fēng)格遷移時(shí)，我們不斷抽取合成圖像的內(nèi)容特征和風(fēng)格特征，然后計(jì)算損失函數(shù)。下面定義了訓(xùn)練循環(huán)。

def train(X, contents_Y, styles_Y, device, lr, num_epochs, lr_decay_epoch):
    X, styles_Y_gram, trainer = get_inits(X, device, lr, styles_Y)
    scheduler = torch.optim.lr_scheduler.StepLR(trainer, lr_decay_epoch, 0.8)
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[10, num_epochs],
                            legend=['content', 'style', 'TV'],
                            ncols=2, figsize=(7, 2.5))
    for epoch in range(num_epochs):
        trainer.zero_grad()
        contents_Y_hat, styles_Y_hat = extract_features(
            X, content_layers, style_layers)
        contents_l, styles_l, tv_l, l = compute_loss(
            X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram)
        l.backward()
        trainer.step()
        scheduler.step()
        if (epoch + 1) % 10 == 0:
            animator.axes[1].imshow(postprocess(X))
            animator.add(epoch + 1, [float(sum(contents_l)),
                                     float(sum(styles_l)), float(tv_l)])
    return X

現(xiàn)在我們訓(xùn)練模型： 首先將內(nèi)容圖像和風(fēng)格圖像的高和寬分別調(diào)整為300和450像素，用內(nèi)容圖像（train函數(shù)中的第一個(gè)參數(shù)）來(lái)初始化合成圖像。

device, image_shape = d2l.try_gpu(), (300, 450)
net = net.to(device)
content_X, contents_Y = get_contents(image_shape, device)
_, styles_Y = get_styles(image_shape, device)
output = train(content_X, contents_Y, styles_Y, device, 0.3, 500, 50)

三種損失的比重也是可以進(jìn)行訓(xùn)練的。

如果想要更好地模仿出圖片風(fēng)格，可以嘗試使用更大的輸出分辨率。

當(dāng)遇到大圖片難以訓(xùn)練的問(wèn)題時(shí)，可以先用較小的分辨率（例如300*450）進(jìn)行訓(xùn)練，將輸出圖像放大后作為初始圖像輸入網(wǎng)絡(luò)進(jìn)行大分辨率輸出（如1000*1200）的訓(xùn)練。

知識(shí)補(bǔ)充

·協(xié)方差矩陣和格萊姆矩陣是計(jì)算的二階統(tǒng)計(jì)信息。統(tǒng)計(jì)直方圖能夠計(jì)算更高階的統(tǒng)計(jì)信息，但是在樣式遷移問(wèn)題中，只需要匹配一階和二階信息就足夠了。

·TV損失可以理解為一種圖像平滑技術(shù)

·GAN在樣式遷移任務(wù)中表現(xiàn)更好