INN实现理解-gaussian_INN

github 地址:https://github.com/hagabbar/cINNamon

首先,这是通过 PyTorch 实现的。PyTorch 和 TensorFlow 一样,都是深度学习框架。

一、PyTorch

教程:PyTorch Handbook

该教程中涉及的内容很多,由于时间有限,暂时就只学习了基础部分,有需要再进行深入学习。

PyTorch 是什么

Torch 是一个与 Numpy 类似的张量(Tensor)操作库,与 Numpy 不同的是 Torch 对 GPU 支持的很好,Lua 是 Torch 的上层包装。PyTorch 和 Torch 使用相同的 C 库,即 PyTorch 和 Torch 使用相同的底层,只是使用了不同的上层包装语言。一句话总结:PyTorch 是一个基于 Torch 的 Python 开源机器学习库,用于自然语言处理等应用程序。

PyTorch是一个Python包,提供两个高级功能:

  • 具有强大的GPU加速的张量计算(如NumPy)
  • 包含自动求导系统的的深度神经网络

基于Python的科学计算包,服务于以下两种场景:

  • 作为NumPy的替代品,可以使用GPU的强大计算能力
  • 提供最大的灵活性和高速的深度学习研究平台
PyTorch v.s. TensorFlow

结论:PyTorch 更有利于研究人员、爱好者、小规模项目等快速搞出原型。而 TensorFlow 更适合大规模部署,特别是需要跨平台和嵌入式部署时。

详细分析见:PyTorch还是TensorFlow?这有一份新手深度学习框架选择指南

Tensor 张量

张量的操作就不贴出来了,不懂的可以查看操作教程.

张量是 PyTorch 里面基础的运算单位,与 Numpy 的 ndarray 相同,都表示的是一个多维的矩阵。与 ndarray 的最大区别就是,PyTorch 的 Tensor 可以在 GPU 上运行,而 numpy 的 ndarray 只能在 CPU 上运行,在 GPU 上运行大大加快了运算速度。

在同构的意义下,第零阶张量(r = 0)为标量(Scalar),第一阶张量(r = 1)为向量(Vector),第二阶张量(r = 2)则成为矩阵(Matrix),第三阶以上的统称为多维张量。

将 Torch Tensor 转换为 NumPy 数组是一件轻而易举的事(反之亦然)。Torch Tensor 和 NumPy 数组将共享底层内存位置,更改一个将改变另一个。

Torch Tensor → NumPy Array

NumPy Array → Torch Tensor

CUDA 张量

Autograd 自动求导机制

PyTorch 中所有神经网络的核心是 autograd 包,autograd 包为张量上的所有操作提供了自动求导。 它是一个在运行时定义的框架,这意味着反向传播是根据你的代码来确定如何运行,并且每次迭代可以是不同的。

torch.Tensor 是这个包的核心类。如果设置 .requires_grad 为 True,那么将会追踪所有对于该张量的操作。.requires_grad_( ... ) 可以改变现有张量的 requires_grad属性,如果没有指定的话,默认输入的 flag 是 False。当完成计算后通过调用 .backward(),自动计算所有的梯度, 这个张量的所有梯度将会自动积累到 .grad 属性。要阻止张量跟踪历史记录,可以调用 .detach() 方法将其与计算历史记录分离,并禁止跟踪它将来的计算记录。为了防止跟踪历史记录(和使用内存),可以将代码块包装在 with torch.no_grad():中。这在评估模型时特别有用,因为模型可能具有 requires_grad = True 的可训练参数,但是我们不需要梯度计算。

在自动梯度计算中还有另外一个重要的类 Function。Tensor 和 Function 互相连接并生成一个非循环图,它表示和存储了完整的计算历史。每个张量都有一个 .grad_fn 属性,这个属性引用了一个创建了 Tensor 的 Function(除非这个张量是用户手动创建的,即,这个张量的 grad_fn 是 None)。

如果需要计算导数,可以在 Tensor 上调用 .backward()。 如果 Tensor 是一个标量(即它包含一个元素数据)则不需要为 backward() 指定任何参数,但是如果它有更多的元素,你需要指定一个 gradient 参数来匹配张量的形状。

可以用 autograd 做更多操作:

如果 .requires_grad=True 但是你又不希望进行 autograd 的计算, 那么可以将变量包裹在 with torch.no_grad() 中。这个方法在测试集测试准确率的时候会经常用到。如:

神经网络

PyTorch 使用 torch.nn 包来构建神经网络。nn 包依赖 autograd 包来定义模型并求导。一个 nn.Module 包含各个层和一个 forward(input) 方法,该方法返回 output。

神经网络反向传播更新网络的参数,主要使用如下简单的更新原则:

weight = weight - learning_rate * gradient

以下图所示的网络结构为例,使用 PyTorch 训练一个网络模型:

  1. 定义网络

在模型中必须要定义 forward 函数,backward 函数(用来计算梯度)会被autograd自动创建。可以在 forward 函数中使用任何针对 Tensor 的操作。

一些常用的操作:

  1. 损失函数

一个损失函数接受一对 (output, target) 作为输入,计算一个值来估计网络的输出和目标值相差多少。nn 包中有很多不同的损失函数, nn.MSELoss 是一个比较简单的损失函数,它计算输出和目标间的均方误差, 例如:

得到 loss 这一过程经历了如下计算:

  1. 反向传播

调用 loss.backward() 进行反向传播。但是在调用前需要清除已存在的梯度,否则梯度将被累加到已存在的梯度。

注意:torch.nn 只支持 mini-batches,不支持一次只输入一个样本,即一次必须是一个 batch。也就是说,就算我们输入一个样本,也会对样本进行分批,所以,所有的输入都会增加一个维度,用于记录 batch-size。

  1. 更新权重

在实践中最简单的权重更新规则是随机梯度下降(SGD):

weight = weight - learning_rate * gradient

我们可以使用简单的Python代码实现这个规则:

learning_rate = 0.01
    for f in net.parameters():
        f.data.sub_(f.grad.data * learning_rate)

但是当使用神经网络是想要使用各种不同的更新规则时,比如 SGD、Nesterov-SGD、Adam、RMSPROP 等,PyTorch 中构建了一个包 torch.optim 实现了所有的这些规则。使用它们非常简单:

二、INN 实现理解

从 examples / gaussian_INN 的 main 开始理解。

  1. 设置超参数:
  1. 生成数据

^^ pos 是一个 (tot_dataset_size, 2) 的矩阵,其元素都是 [0, 1) 区间的浮点数;labels 是一个 (tot_dataset_size, ndata) 的矩阵;x 是一个大小为 ndata 的向量,其第 i 个元素的值为 (i-1) / float(ndata);sig 也是一个 (tot_dataset_size, ndata) 的矩阵,其与 data 的关系为 sig + noise = labels,noise 是一个 (tot_dataset_size, ndata) 的矩阵,其元素服从正态分布 N(0, sigma)。

然后分别取 pos、labels、sig 的倒数 test_split 个元素作为测试数据,画出测试数据的分布图如下:

^^ 每幅图对应一个样本,图中蓝点为 labels 数据,红线为 sig 数据。

  1. 建立模型

^^ 这里有个疑问:为什么 latent z 的维度是指定的而非根据 x 和 y 的维度计算出来的?

^^ 输入结点的维度数应为 x 的显示维度,但当 y + z 的显示维度大于 x 的显式维度时,需要对 x 用 0 补齐。因此输入结点的维度为 ndim_tot。

^^ 三层隐藏层

^^ ReversibleGraphNet 类表示了可逆网络,它是 torch.nn.Module 类的子类

  1. 训练前准备工作

^^ 训练参数

^^ 各项损失的相对权重。INN 训练过程中考虑三项损失:①模型输出 yi = s(xi) 与网络预测 fy(xi) 之间的偏差,损失记为 Ly(yi,fy(xi)),Ly 可以是任意有监督的损失;lamdb_predict 为 Ly 的权重;②模型输出 p(y = s(x)) = p(x) / |Js| 和潜在变量 p(z) 的边际分布的乘积与网络输出 q(y = fy(x),z = fz(x)) = p(x) / |Jyz| 间的偏差,记为 Lz(p(y)p(z),q(y,z));lambd_latent 为 Lz 的权重;③输入端的损失 LxLx(p(x),q(x)) 表示了 p(x) 与后向预测分布 q(x) = p(y = fy(x)) p(z = fz(x)) / |Jx| 间的偏差;lambd_rev 为 Lx 的权重.

^^ 定义权重更新规则(optimizer),scheduler 对其进行封装,目的是使其学习率每隔 step_size 轮就进行一次衰减。

^^ 定义损失函数。需要定义三个函数来进行三种损失的计算,其中 Lx、Lz 是无监督损失,因此选择了 MMD_multiscale(多刻度的 MMD,MMD 常用于度量两个不同但相关的分布的距离)作为损失函数;Ly 是有监督损失,因此选择了 fit(即平方误差)。

^^ 建立训练集数据装载器。DataLoader 返回的是一个迭代器,我们可以使用迭代器分批获取数据,或者直接使用 for 循环对其进行遍历。

^^ 选取训练数据(这里只有 3*3 个样本),预先计算其真实概率这里没看懂

  1. 训练模型

这个实现对 INN 训练了 12000 次。我们只看一次训练的步骤。

首先是设置此轮训练的学习率,这个一般由我们之前包装的 scheduler 根据训练轮数来进行设定。

^^ 训练网络。其中调用了 train() 函数:


核心函数 train()

首先要将训练涉及的各个模块的状态设置为 training。

^^ 设置 loss_factor,当 i_epoch 大于 300 时,其值为 1。

每轮训练只能使用数据装载器装载 n_its_per_epoch(设定为12)批的样本数据,对于每一批数据,进行如下处理。

^^ 对 x 和 yz 进行对其填充,使它们的维度和 ndim_tot 相同。在填充前,先为 y 增加随机噪声。

^^ 在开始训练前,需要清除已存在的梯度。

^^ 执行正向传播得到输出 output(正向计算由 PyTorch 实现),output 与输入有相同的维度(这里是 128328,其中,128200 记录得到的 z',128*128 记录得到的 y')。然后对 y 进行修改,y_short 等于正向传播前的 y 去掉中间对齐填充的部分(这里由于 dim(y) + dim(z) = tot_dim,因此没有进行对其填充,即 y_short 即 z 与 y 的连接 )。

^^ 计算损失 Ly,即为 y 和 y' 的均方误差,记录在 l_forward 中。

^^ 计算损失 Lz

^^ 这个 backward() 函数是 PyTorch 实现的,调用它是为了获得反向传播的误差。


^^ 将此轮的损失值追加到历史记录中,用于之后的可视化。

之后,每隔 50 轮训练,就将训练结果可视化出来(包括可视化 latent-space 和测试结果),这里涉及到很多可视化内容,暂时就先不看了。

推荐阅读更多精彩内容