【生成模型】简述概率密度函数可处理流模型

本期将介绍第二种非常优雅的生成模型—流模型，它也是一种概率密度函数可处理的生成模型。本文将对其原理进行介绍，并对nice模型的源码进行讲解。

                                          作者&编辑 | 小米粥

1 流模型

这是一种想法比较直接但实际不容易构造的生成模型，它通过可逆的非线性变换等技巧使得似然函数可以被精确计算出来。对于分布比较简单（例如高斯分布）的隐变量z，其概率分布为 p_z(z) ，这时若存在一个连续、可微、可逆的非线性变换g(z)，将简单的潜变量z的分布转换成关于样本x的一个复杂分布，将非线性变换g(z)的逆变换记为f(x)，则可得到样本x的准确的概率密度函数 p_x(x)

注意，非线性变换g(z)会引起空间的变形，即p_x(x)不等于 p_z(f(x))，且有p_z(z)dz=p_x(x)dx。若上述模型构建成功，则生成样本时只需从简单分布p_z(z)中随机采样然后将其变换为 x=g(z) 即可。

为了训练非线性独立成分估计模型，我们必须计算样本的概率密度函数p_x(x)。分析上式，概率密度函数p_x(x)的计算需要计算p_z(z)和雅可比矩阵的行列式绝对值。对于前者，要构建一个 g(z) 的逆变换f(x)，这样当给定一个样本 x 后，就可以通过z=f(x)计算得到其对应的隐变量，p_z(z)通常设计为简单的分布，容易计算；对于后者，要将f(x)设计为某种特殊的形式，使得雅可比矩阵的行列式易于计算。另外，变换的可逆性要求样本x和隐变量z具有相同的维度。综上，需要将生成模型精心设计成一种易于处理且灵活的双射模型，使其逆变换f(x)存在，且对应的雅可比矩阵的行列式可计算。

为了对这类模型有更深刻的理解，我们对NICE模型进行详细的介绍。NICE模型的逆变换f(x)由多个加性耦合层和一个尺度变换层构成。在每个加性耦合层中，首先将 n 维样本 x 分解为两部分x₁和x₂，例如将x的第1,3,5...个元素划入x₁部分，将第2,4,6...个元素划入x₂部分，每个部分的维度均为n/2；也可以将 x 的第2,4,6...个元素划入x₁部分，将第1,3,5...个元素划入 x₂ 部分；或者使用其他划分方式。然后对两部分其进行变换：

其中m()为任意函数，注意这里要保证m()的输出结果维度与 x₂ 保持一致，NICE模型使用多层全连接网络和ReLU激活函数来构建 m() 。容易发现，使用加性耦合层作为逆变换f(x)的其中一部分，它是可逆的并且雅可比矩阵的行列式也是容易计算的。当已知h₁和h₂时，可得其逆变换:

其雅可比矩阵为：

根据三角矩阵的性质，其行列式为对角元素的乘积，故加性耦合层雅可比矩阵的行列式绝对值为1。当将多个加性耦合层串联时，如下图所示。由于每个层的逆变换是容易计算的，则串联后的逆变换仍然是容易计算的。此时的雅可比矩阵为:

根据矩阵行列式的性质，有：

需要说明的是，必须注意在不同的加性耦合层使用不同的划分策略，使得样本不同维度的信息充分混淆。在尺度变换层，定义了包含n个非负参数的向量s=[s₁,s₂,...,s_n]，将加性耦合层的输出结果h^(l)与s逐元素相乘可得到对应的隐变量z。这里s用于控制每个维度的特征变换的尺度，可以表征维度的重要性，对应维度的数值较大表明这一维度的重要性低，因为生成样本时隐变量需要先经过尺度变换层，隐变量在尺度变换层需要逐元素乘1/s。显然，尺度变换层的逆变换只需逐元素乘1/s，为了计算雅可比矩阵，将尺度变换写为对角矩阵的形式：

则其雅可比矩阵的行列式为s₁s₂...s_n。现在，我们构造了可逆的、雅可比矩阵的行列式绝对值易于计算的逆变换f(x)，对于隐变量z，NICE模型假设其n个维度彼此独立，即

若选择z为高斯分布，则样本x的似然函数为：

若选择z为logistic分布，即

则样本x的似然函数为

现在，我们可以使用极大似然法对NICE模型进行训练，训练完成后也得到了生成模型g(z)。若  z 为高斯分布，则直接从高斯分布中采样可得到z；若选择  为logistic分布，可现在0-1之间的均匀分布中采样得到\epsilon ，然后使用变换z=t(ε)得到隐变量。根据两个随机变量的映射关系

则有

将隐变量z使用非线性变换g(z)，即经过尺度变换层的逆变换、多个加性耦合层的逆变换可得到生成样本x。

2 NICE 代码

接下来我们将提供一份完整的NICE的代码讲解，其中训练集为mnist数据集。

首先读取相关python库：

import argparse 

import torch 

import torchvision 

import numpy as np 

import niceimport utils

定义主函数，设置相关参数，并进行模型的训练和验证

定义二维掩膜卷积核，其中有A与B两种类型，区别之处在于中心位置是否被卷积计算：

def main(args):

    device = torch.device("cuda:0") 

    # 训练相关参数设置

    dataset = args.dataset

    batch_size = args.batch_size

    latent = args.latent

    max_iter = args.max_iter

    sample_size = args.sample_size

    coupling = 4    mask_config = 1.

    # 优化器设置

    lr = args.lr

    momentum = args.momentum

    decay = args.decay

    zca = None

    mean = None

    # 选择相应的训练数据集

    if dataset == 'mnist':

        mean = torch.load('./statistics/mnist_mean.pt')        (full_dim, mid_dim, hidden) = (1 * 28 * 28, 1000, 5)        transform = torchvision.transforms.ToTensor()        trainset = torchvision.datasets.MNIST(root='~/torch/data/MNIST',            train=True, download=True, transform=transform)        trainloader = torch.utils.data.DataLoader(trainset,            batch_size=batch_size, shuffle=True, num_workers=2)         

  # 使用正态分布或逻辑回归分布

   if latent == 'normal':

        prior = torch.distributions.Normal(

            torch.tensor(0.).to(device), torch.tensor(1.).to(device))

    elif latent == 'logistic':

        prior = utils.StandardLogistic()

    filename = '%s_' % dataset \

             + 'bs%d_' % batch_size \

             + '%s_' % latent \

             + 'cp%d_' % coupling \

             + 'md%d_' % mid_dim \

             + 'hd%d_' % hidden

    # 实例化流模型

    flow = nice.NICE(prior=prior,

                coupling=coupling,

                in_out_dim=full_dim,

                mid_dim=mid_dim,

                hidden=hidden,

                mask_config=mask_config).to(device)

    optimizer = torch.optim.Adam(

        flow.parameters(), lr=lr, betas=(momentum, decay), eps=1e-4)

    total_iter = 0

    train = True

    running_loss = 0

    # 训练模型

    while train:

        for _, data in enumerate(trainloader, 1):

            flow.train()

    # set to training mode

            if total_iter == max_iter:

                train = False

                break

            total_iter += 1

            optimizer.zero_grad()

    # clear gradient tensors

            inputs, _ = data

            inputs = utils.prepare_data(

                inputs, dataset, zca=zca,

 mean=mean).to(device)

            # log-likelihood of input minibatch

            loss = -flow(inputs).mean()

            running_loss += float(loss)

            # backprop and update parameters

            loss.backward()

            optimizer.step()

            if total_iter % 10 == 0:

                mean_loss = running_loss / 2000

                bit_per_dim = (mean_loss + np.log(256.) * full_dim) \

                            / (full_dim * np.log(2.)) 

                print('iter %s:' % total_iter,

                    'loss = %.3f' % mean_loss,

                    'bits/dim = %.3f' % bit_per_dim)

                running_loss = 0.0

            if total_iter % 1000 == 0:

                flow.eval()

        # set to inference mode

                with torch.no_grad():

                    samples = flow.sample(sample_size).cpu()

                    samples = utils.prepare_data(                        samples, dataset, zca=zca, mean=mean, reverse=True)

                torchvision.utils.save_image(torchvision.utils.make_grid(samples),

'./samples/' + filename +'iter%d.png' % total_iter)

    # 保存模型

    torch.save({

        'total_iter': total_iter,

        'model_state_dict': flow.state_dict(),

        'optimizer_state_dict': optimizer.state_dict(),

        'dataset': dataset,

        'batch_size': batch_size,

        'latent': latent,

        'coupling': coupling,

        'mid_dim': mid_dim,

        'hidden': hidden,

        'mask_config': mask_config},

        './models/mnist/' + filename +'iter%d.tar' % total_iter)

    print('Checkpoint Saved')

 if __name__ == '__main__':

    parser = argparse.ArgumentParser('MNIST NICE PyTorch implementation')

    parser.add_argument('--dataset',

                        help='dataset to be modeled.',

                        type=str,

                        default='mnist')    parser.add_argument('--batch_size', 

                        help='number of images in a mini-batch.',

                        type=int,

                        default=256)

    parser.add_argument('--latent',

                        help='latent distribution.',

                        type=str,

                        default='logistic')

    parser.add_argument('--max_iter',

                        help='maximum number of iterations.',

                        type=int,

                        default=20000)

    parser.add_argument('--sample_size',

                        help='number of images to generate.',

                        type=int,

                        default=64)

    parser.add_argument('--lr',

                        help='initial learning rate.',

                        type=float,

                        default=1e-3)

    parser.add_argument('--momentum', 

                        help='beta1 in Adam optimizer.',

                        type=float,

                        default=0.9)

    parser.add_argument('--decay',

                        help='beta2 in Adam optimizer.',

                        type=float,

                        default=0.999)

    args = parser.parse_args()

    main(args)

NICE流模型的核心代码：

# 加性耦合层class Coupling(nn.Module):

    def __init__(self, in_out_dim, mid_dim, hidden, mask_config):

        super(Coupling, self).__init__()

        self.mask_config = mask_config

        self.in_block = nn.Sequential(

            nn.Linear(in_out_dim//2, mid_dim),

            nn.ReLU())

        self.mid_block = nn.ModuleList([

            nn.Sequential(

                nn.Linear(mid_dim, mid_dim),

                nn.ReLU()) for _ in range(hidden - 1)])

        self.out_block = nn.Linear(mid_dim, in_out_dim//2)

    def forward(self, x, reverse=False):

        [B, W] = list(x.size())

        x = x.reshape((B, W//2, 2))

        if self.mask_config:

            on, off = x[:, :, 0], x[:, :, 1]

        else:

            off, on = x[:, :, 0], x[:, :, 1]

        off_ = self.in_block(off)

        for i in range(len(self.mid_block)):

            off_ = self.mid_block[i](off_)

        shift = self.out_block(off_)

        if reverse:

            on = on - shift

        else:

            on = on + shift

        if self.mask_config:

            x = torch.stack((on, off), dim=2)

        else:

            x = torch.stack((off, on), dim=2)

        return x.reshape((B, W))

 # 尺度变换层

 class Scaling(nn.Module):

    def __init__(self, dim):

        super(Scaling, self).__init__()

        self.scale = nn.Parameter(

            torch.zeros((1, dim)), requires_grad=True)

    def forward(self, x, reverse=False):

        log_det_J = torch.sum(self.scale)

        if reverse:

            x = x * torch.exp(-self.scale)

        else:

            x = x * torch.exp(self.scale)

        return x, log_det_J

 # NICE模型

 class NICE(nn.Module):

    def __init__(self, prior, coupling,

        in_out_dim, mid_dim, hidden, mask_config):

        self.prior = prior

        self.in_out_dim = in_out_dim

        self.coupling = nn.ModuleList([

            Coupling(in_out_dim=in_out_dim,

                     mid_dim=mid_dim,

                     hidden=hidden,

                     mask_config=(mask_config+i)%2) \            for i in range(coupling)])

        self.scaling = Scaling(in_out_dim)

    def g(self, z):

        x, _ = self.scaling(z, reverse=True)

        for i in reversed(range(len(self.coupling))):

            x = self.coupling[i](x, reverse=True)

        return x

    def f(self, x):

        for i in range(len(self.coupling)):

            x = self.coupling[i](x)

        return self.scaling(x)

    def log_prob(self, x):

        z, log_det_J = self.f(x)

        log_ll = torch.sum(self.prior.log_prob(z), dim=1)

        return log_ll + log_det_J

    def sample(self, size):

        z = self.prior.sample((size, self.in_out_dim)).cuda()        return self.g(z)

    def forward(self, x):

        return self.log_prob(x)

[1] Dinh L , Krueger D , Bengio Y . NICE: Non-linear Independent Components Estimation[J]. Computer ence, 2014.

[2]Diederik P. Kingma*†, Prafulla Dhariwal⇤. Glow: Generative Flow with Invertible 1x1 Convolutions[J]. 2018.

总结

本期带大家学习了流模型，流模型的生成效果也非常卓越，但是其网络设计比较复杂，训练难度比较大。下一期将为大家继续讲解其他生成模型。

个人知乎，欢迎关注

GAN群

有三AI建立了一个GAN群，便于有志者相互交流。感兴趣的同学也可以微信搜索xiaozhouguo94，备注“加入有三-GAN群”。

更多GAN的学习

知识星球是有三AI的付费内容社区，里面包超过100种经典GAN模型的解读，了解详细请阅读以下文章：

【杂谈】有三AI知识星球指导手册出炉！和公众号相比又有哪些内容？

有三AI秋季划GAN学习小组，可长期跟随有三学习GAN相关的内容，并获得及时指导，了解详细请阅读以下文章：

【杂谈】如何让2020年秋招CV项目能力更加硬核，可深入学习有三秋季划4大领域32个方向

转载文章请后台联系

侵权必究

往期精选

• 【GAN优化】GAN优化专栏上线，首谈生成模型与GAN基础

• 【GAN的优化】从KL和JS散度到fGAN

• 【GAN优化】详解对偶与WGAN

• 【GAN优化】详解SNGAN(频谱归一化GAN)

• 【GAN优化】一览IPM框架下的各种GAN

• 【GAN优化】GAN优化专栏栏主小米粥自述，脚踏实地，莫问前程

• 【GAN优化】GAN训练的几个问题

• 【GAN优化】GAN训练的小技巧

• 【GAN优化】从动力学视角看GAN是一种什么感觉？

• 【GAN优化】小批量判别器如何解决模式崩溃问题

• 【GAN优化】长文综述解读如何定量评价生成对抗网络(GAN)
  

• 【技术综述】有三说GANs（上）
  

• 【模型解读】历数GAN的5大基本结构

• 【百战GAN】如何使用GAN拯救你的低分辨率老照片
  

• 【百战GAN】二次元宅们，给自己做一个专属动漫头像可好！
  

• 【百战GAN】羡慕别人的美妆？那就用GAN复制粘贴过来
  

• 【百战GAN】GAN也可以拿来做图像分割，看起来效果还不错？
  

• 【百战GAN】新手如何开始你的第一个生成对抗网络(GAN)任务

• 【百战GAN】自动增强图像对比度和颜色美感，GAN如何做？

• 【直播回放】80分钟剖析GAN如何从各个方向提升图像的质量

• 【直播回放】60分钟剖析GAN如何用于人脸的各种算法