[nflow学习] demo学习

demo 1

依赖库安装：

pip install nflows

import matplotlib.pyplot as plt
import sklearn.datasets as datasets

import torch
from torch import nn
from torch import optim

from nflows.flows.base import Flow
from nflows.distributions.normal import StandardNormal
from nflows.transforms.base import CompositeTransform
from nflows.transforms.autoregressive import MaskedAffineAutoregressiveTransform
from nflows.transforms.permutations import ReversePermutation

data

x, y = datasets.make_moons(128, noise=.1)
plt.scatter(x[:, 0], x[:, 1])
plt.show()

其中x.shape=[128, 2] y.shape=[128]

▲ sklearn中datasets对象创建数据集的几种方式（感兴趣的可以自己尝试画图看下)：

网络搭建

num_layers = 5
base_dist = StandardNormal(shape=[2])

transforms = []
for _ in range(num_layers):
    transforms.append(ReversePermutation(features=2))
    transforms.append(MaskedAffineAutoregressiveTransform(features=2, 
                                                          hidden_features=4))
transform = CompositeTransform(transforms)

flow = Flow(transform, base_dist)
optimizer = optim.Adam(flow.parameters())

网络一共5层，每层由一个ReversePermutation和MaskedAffineAutoregressiveTransform组成。
transform为flow模型，base_dist为基础分布（StandardNormal）

模型训练

num_iter = 5000
for i in range(num_iter):
    x, y = datasets.make_moons(128, noise=0.1)
    x = torch.tensor(x, dtype=torch.float32)
    optimizer.zero_grad()
    loss = -flow.log_prob(inputs=x).mean()
    loss.backward()
    optimizer.step()

图像绘制

    if(i+1) % 500 == 0:
        xline = torch.linspace(-1.5, 2.5)
        yline = torch.linspace(-0.75, 1.25)
        xgrid, ygrid = torch.meshgrid(xline, yline)
        xyinput = torch.cat([xgrid.reshape(-1, 1), ygrid.reshape(-1, 1)], dim=1)

        with torch.no_grad():
            zgrid = flow.log_prob(xyinput).exp().reshape(100, 100)

        plt.contourf(xgrid.numpy(), ygrid.numpy(), zgrid.numpy())
        plt.title('iteration {}'.format(i+1))
        plt.show()

结果展示

拟合结果对比

demo2

框架与demo1相同，部分有变化：

加了参数context:

base_dist = ConditionalDiagonalNormal(shape=[2], context_encoder=nn.Linear(1,4))

transforms.append(MaskedAffineAutoregressiveTransform(features=2, hidden_features=4, context_features=1))

y = torch.tensor(y, dtype=torch.float32).reshape(-1, 1)

loss = -flow.log_prob(inputs=x, context=y).mean()

结果展示

拟合结果对比

结论

从图像直观感受，demo2的拟合结果是比demo1相对好一些。这与网络和参数变化有关。