风格迁移训练实践与分析

前一篇文章分享了Pytorch简单风格迁移的代码，本着不跑挂服务器不死心的态度，不停的增加计算步骤，看看图片融合生成的效果，

为了方便一次性执行，把代码简单改造了一下，与前一篇文章大同小异：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from PIL import Image
import matplotlib.pyplot as plt

import torchvision.transforms as transforms
import torchvision.models as models
import datetime

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def get_img_size(img_name):
    """
    获取图像大小
    :param img_name:
    :return:
    """
    im = Image.open(img_name).convert('RGB')
    return im, im.height, im.width


def image_loader(img, im_h, im_w):
    """
    加载图像
    :param img:
    :param im_h:
    :param im_w:
    :return:
    """

    # loader = transforms.Compose([transforms.Resize([im_h, im_w]), transforms.ToTensor()])
    loader = transforms.Compose([transforms.Resize([1000, 1000]), transforms.ToTensor()])
    im_l = loader(img).unsqueeze(0)
    return im_l.to(device, torch.float)


def im_show(tensor, save_file_path):
    """
    显示保存图片
    :param tensor:
    :param save_file_path:
    :return:
    """
    image = tensor.cpu().clone()
    image = image.squeeze(0)
    image = transforms.ToPILImage()(image)
    plt.imshow(image, aspect='equal')
    plt.axis('off')
    plt.savefig(save_file_path, bbox_inches='tight', pad_inches=0.0)
    plt.pause(0.001)


class ContentLoss(nn.Module):
    """
    内容损失
    """

    def __init__(self, target,):
        super(ContentLoss, self).__init__()
        self.target = target.detach()

    def forward(self, cl_input):
        self.loss = F.mse_loss(cl_input, self.target)
        return cl_input


def gram_matrix(gm_input):
    """
    风格损失矩阵
    :param gm_input:
    :return:
    """
    a, b, c, d = gm_input.size()
    features = gm_input.view(a * b, c * d)
    G = torch.mm(features, features.t())

    return G.div(a * b * c * d)


class StyleLoss(nn.Module):
    """
    风格损失
    """

    def __init__(self, target_feature):
        super(StyleLoss, self).__init__()
        self.target = gram_matrix(target_feature).detach()

    def forward(self, fw_input):
        G = gram_matrix(fw_input)
        self.loss = F.mse_loss(G, self.target)
        return fw_input


# 使用19层的VGG神经网络模型
cnn = models.vgg19(pretrained=True).features.to(device).eval()


cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)


class Normalization(nn.Module):
    """
    规范化输入图像
    """
    def __init__(self, mean, std):
        super(Normalization, self).__init__()
        self.mean = mean.clone().detach().view(-1, 1, 1)
        self.std = std.clone().detach().view(-1, 1, 1)

    def forward(self, img):
        return (img - self.mean) / self.std


def get_style_model_and_losses(cn, normalization_mean, normalization_std, style_i, content_i, cld, sld):
    """
    获取内容损失和风格损失
    :param cn:
    :param normalization_mean:
    :param normalization_std:
    :param style_i:
    :param content_i:
    :param cld:
    :param sld:
    :return:
    """

    normalization = Normalization(normalization_mean, normalization_std).to(device)
    content_losses = []
    style_losses = []

    model = nn.Sequential(normalization)

    i = 0
    for layer in cn.children():
        if isinstance(layer, nn.Conv2d):
            i += 1
            name = 'conv_{}'.format(i)
        elif isinstance(layer, nn.ReLU):
            name = 'relu_{}'.format(i)
            layer = nn.ReLU(inplace=False)
        elif isinstance(layer, nn.MaxPool2d):
            name = 'pool_{}'.format(i)
        elif isinstance(layer, nn.BatchNorm2d):
            name = 'bn_{}'.format(i)
        else:
            raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__))

        model.add_module(name, layer)

        if name in cld:
            target = model(content_i).detach()
            content_loss = ContentLoss(target)
            model.add_module("content_loss_{}".format(i), content_loss)
            content_losses.append(content_loss)

        if name in sld:
            target_feature = model(style_i).detach()
            style_loss = StyleLoss(target_feature)
            model.add_module("style_loss_{}".format(i), style_loss)
            style_losses.append(style_loss)

    for i in range(len(model) - 1, -1, -1):
        if isinstance(model[i], ContentLoss) or isinstance(model[i], StyleLoss):
            break

    model = model[:(i + 1)]

    return model, style_losses, content_losses


def get_input_optimizer(input_i):
    """
    使用 L-BFGS 算法
    最小化风格、内容的损失
    :param input_i:
    :return:
    """
    optimizer = optim.LBFGS([input_i])
    return optimizer


def run_style_transfer(cn, norma_mean, normalization_std, ct_img, sl_img, in_img, steps, style_weight, content_weight):
    """
    样式转换，建立风格迁移模型
    :param cn:
    :param norma_mean:
    :param normalization_std:
    :param ct_img:
    :param sl_img:
    :param in_img:
    :param steps:
    :param style_weight:
    :param content_weight:
    :return:
    """
    content_layers = ['conv_4']
    style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
    model, style_losses, content_losses = get_style_model_and_losses(cn, norma_mean, normalization_std, sl_img, ct_img, content_layers, style_layers)
    in_img.requires_grad_(True)
    model.requires_grad_(False)

    optimizer = get_input_optimizer(in_img)
    print('Optimizing..')
    run = [0]
    while run[0] <= steps:

        def closure():
            with torch.no_grad():
                in_img.clamp_(0, 1)

            optimizer.zero_grad()
            model(in_img)
            style_score = 0
            content_score = 0

            for sl in style_losses:
                style_score += sl.loss
            for cl in content_losses:
                content_score += cl.loss

            style_score *= style_weight
            content_score *= content_weight

            loss = style_score + content_score
            loss.backward()

            run[0] += 1
            if run[0] % 50 == 0:
                print("run {}:".format(run))
                print('Style Loss : {:4f} Content Loss: {:4f}'.format(style_score.item(), content_score.item()))
            return style_score + content_score

        optimizer.step(closure)
    with torch.no_grad():
        in_img.clamp_(0, 1)
    return in_img


def style_transfer(content_image_path, style_image_path, image_save_path, run_steps):
    """
    风格迁移主入口
    :param content_image_path: 内容图片
    :param style_image_path: 风格图片
    :param image_save_path: 存储图片地址
    :param run_steps: 执行计算次数
    :return:
    """
    c_image, c_im_h, c_im_w = get_img_size(content_image_path)
    s_image, s_im_h, s_im_w = get_img_size(style_image_path)
    content_img = image_loader(c_image, c_im_h, c_im_w)
    style_img = image_loader(s_image, c_im_h, c_im_w)
    assert style_img.size() == content_img.size()
    # 输入内容图像
    input_img = content_img.clone()
    begin_time = datetime.datetime.now()
    print("******************开始时间*****************", begin_time)
    output = run_style_transfer(cnn, cnn_normalization_mean, cnn_normalization_std, content_img, style_img, input_img, run_steps, s_weight, c_weight)
    try:
        im_show(output, image_save_path)
    except Exception as e:
        print(e)
    print("******************结束时间*****************", datetime.datetime.now())
    print("******************耗时*****************", datetime.datetime.now() - begin_time)


if __name__ == '__main__':
    s_weight = 1000000
    c_weight = 1
    # content_img_path = "data/drew/img/512.png"
    content_img_path = "/data/drew/img/dancing.jpg"
    # style_img_path = "data/drew/img/512r.png"
    style_img_path = "/data/drew/img/picasso.jpg"
    for steps in range(100, 3200, 200):
        # save_path = "data/drew/img/end_%s_%s.jpg" % (steps, datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
        save_path = "/data/drew/img/end_%s_%s.jpg" % (steps, datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
        style_transfer(content_img_path, style_img_path, save_path, steps)

计算步数：

100  耗时16秒

  

 500  耗时1分18秒

  

 900   耗时2分20秒

  

 1300   耗时3分22秒

  

 1700  耗时4分24秒

  

1900    耗时4分54秒

  

2100  耗时5分25秒

  

2300    耗时5分55秒

  

 2900   耗时7分29秒

  

 3100   耗时8分钟

  

5000   耗时12分钟49秒

  

 10000   耗时25分钟33秒

   

聚集对比

 

  

 可以看出，随着计算步骤的增加，图片融合细腻度越来越好，但指不定某个期间取到不同的色值，融合出的图有可能出现较大差异，程序是不知道那张图适合你的，得自己多次计算看那张相对优秀，

个人感觉1700和3100看上去效果都不错，有时候算的越多反而越不适合。总的来说，训练越久，图片的色彩，颜色的深浅，颜色的层次及位置，都越来越与风格图片类似，比如30000次训练的结果，

可以看出越来越接近风格图片。如果一直训练下去，比如训练个百万次，会是啥样？~~O(∩_∩)O，谁的机子牛逼的可以试试！！

 

 计算步数与耗时：

100--------16.042588

300--------45.938147

500--------01:18.010645

700--------01:49.314894

900--------02:20.706951

1100-------02:51.681638

1300-------03:22.309657

1500-------03:53.468090

1700-------04:24.346204

1900-------04:54.553869

2100-------05:25.256131

2300-------05:55.873070

2500-------06:26.832513

2700-------06:57.715941

2900-------07:29.032707

3100-------07:59.818595

5000-------12:49.948595

10000-----25:33.300846

30000-----1:16:40.921489

当然，调整两张图里的权重，迁移融合的效果差异也很大：

比如我们把内容的权重由1调整为10，100，看看有啥效果：

if __name__ == '__main__':
    s_weight = 1000000
    c_weight = 100

  

  

 左图为百万比一，中间图为十万比一，右图为一万比一，看看是不是效果有很大差别。

 例二：

找了一幅竹图和一个竹字，试试效果：

   

结论一样

 

感兴趣的欢迎讨论探索！！