3D VAE神经网络实战

在这个项目中，我们将学习如何使用卷积变分自动编码器 (CVAE) 来处理和重建 3D 湍流数据。

我们使用计算流体动力学 (CFD) 方法生成 3D 湍流立方体，每个 3D 立方体沿着三个速度分量携带物理信息（与图像数据类似，被视为单独的通道）。

作为 3D CFD 数据预处理的一部分，我们编写了一个自定义 pytorch 数据加载器，用于对数据集执行标准化和批处理操作。

CVAE 对预处理后的数据实施 3D 卷积 (3DConvs) 来执行重建。

通过微调超参数和操纵我们的模型架构，我们在 3D 重建方面取得了显着的进步。

可以在这个 Github 存储库中找到项目代码。

1、数据说明

我们的数据集是使用 CFD 模拟方法生成的，它包含从供暖通风和空调 (HVAC) 管道中提取的立方体。

每个立方体代表特定时间携带物理信息的湍流的三维时间快照。从模拟中提取的信息基于两个流动分量：速度场 U 和静压 p。 U 场 (x, y, z) 和标量 p 基于流的方向（立方体的法线方向）。

我们使用体素将 3D 立方体表示为尺寸为 21 × 21 × 21 x 100（x_coord、y_coord、z_coord、timestep）的数组。下图显示了一个立方体数据样本，我们使用热图可视化每个速度分量。

总的来说，该数据集由 96 个模拟组成，每个模拟有 100 个时间步长，总共 9600 个立方体（对于每个速度分量）。

注意：由于保密限制，我们不会公开我们的数据，你可以使用脚本并将其改编为自己的 3D 数据。

2、数据预处理

下面的脚本显示了为预处理 3D 数据而编写的自定义 pytorch 数据加载器。以下是一些亮点：

立方体速度通道的加载和串联
数据标准化
数据缩放

请参阅存储库中的 dataloader.py 脚本以了解完整的实现。

3、模型架构

下图显示了已实现的 CVAE 架构。在本例中，为了清晰起见，显示了 2DConv，但实现的架构使用 3DConv。

CVAE 由编码器网络（上部）、变分层（mu 和 sigma）（中右部分）和解码器网络（底部）组成。

编码器对输入立方体执行下采样操作，解码器对它们进行上采样以恢复原始形状。变分层尝试学习数据集的分布，该层稍后可用于生成。

编码器网络由四个 3D 卷积层组成，每层的卷积滤波器数量是前一层的两倍（分别为 32、64、128 和 256），这使得模型能够学习更复杂的流特征。

密集层用于组合从最后一个编码器层获得的所有特征图，该层连接到计算后验流数据分布的参数（mu和sigma，这些参数使用重新定义概率分布）的变分层。 -[1]中描述的参数化技巧。这种概率分布允许我们从中采样，以生成尺寸为 8 × 8 x 8 的合成 3D 立方体。

解码器网络采用潜在向量并应用四个 3D 转置卷积层来恢复（重建）原始数据维度，每一层的卷积滤波器数量是前一层的一半（分别为 256、128、64 和 32）。

CVAE 使用两个损失函数进行训练：用于重建的均方误差 (MSE) 和用于潜在空间正则化的 Kullback-Leibler 散度 (KLB)。

我们将[2]中提出的架构和[3]中的超参数作为基线架构。

下面的脚本显示了 pytorch 中的一个示例，其中编码器和解码器都是使用 3D 卷积层 (Conv3d) 定义的：

self.encoder = nn.Sequential(
            nn.Conv3d(in_channels=image_channels, out_channels=16, kernel_size=4, stride=1, padding=0),
            nn.BatchNorm3d(num_features=16),
            nn.ReLU(),
            nn.Conv3d(in_channels=16, out_channels=32, kernel_size=4, stride=1, padding=0),
            nn.BatchNorm3d(num_features=32),
            nn.ReLU(),
            nn.Conv3d(in_channels=32, out_channels=64, kernel_size=4, stride=1, padding=0),
            nn.BatchNorm3d(num_features=64),
            nn.ReLU(),
            nn.Conv3d(in_channels=64, out_channels=128, kernel_size=4, stride=1, padding=0),
            nn.BatchNorm3d(num_features=128),
            nn.ReLU(),
            nn.Conv3d(in_channels=128, out_channels=128, kernel_size=4, stride=1, padding=0),
            nn.BatchNorm3d(num_features=128),
            nn.ReLU(),
            Flatten()
        )

self.decoder = nn.Sequential(
            UnFlatten(),
            nn.BatchNorm3d(num_features=128),
            nn.ReLU(),
            nn.ConvTranspose3d(in_channels=128, out_channels=128, kernel_size=4, stride=1, padding=0),
            nn.BatchNorm3d(num_features=128),
            nn.ReLU(),
            nn.ConvTranspose3d(in_channels=128, out_channels=64, kernel_size=4, stride=1, padding=0),
            nn.BatchNorm3d(num_features=64),
            nn.ReLU(),
            nn.ConvTranspose3d(in_channels=64, out_channels=32, kernel_size=4, stride=1, padding=0),
            nn.BatchNorm3d(num_features=32),
            nn.ReLU(),
            nn.ConvTranspose3d(in_channels=32, out_channels=16, kernel_size=4, stride=1, padding=0),
            nn.BatchNorm3d(num_features=16),
            nn.ReLU(),
            nn.ConvTranspose3d(in_channels=16, out_channels=image_channels, kernel_size=4, stride=1, padding=0), # dimensions should be as original
            nn.BatchNorm3d(num_features=3))

4、设置环境

克隆此存储库：

git clone git@github.com:agrija9/Convolutional-VAE-for-3D-Turbulence-Data

建议使用虚拟环境来运行本项目：

可以安装Anaconda并在系统中创建环境
可以使用 pip venv 创建环境

在 pip/conda 环境中安装以下依赖项：

NumPy (>= 1.19.2)
Matplotlib (>= 3.3.2)
PyTorch (>= 1.7.0)
Torchvision (>= 0.8.1)
scikit-learn (>= 0.23.2)
tqdm
tensorboardX
torchsummary
PIL
collections

5、模型训练

要训练模型，请打开终端，激活 pip/conda 环境并输入：

cd /path-to-repo/Convolutional-VAE-for-3D-Turbulence-Data
python main.py --test_every_epochs 3 --batch_size 32 --epochs 40 --h_dim 128 --z_dim 64

以下是一些可以修改的超参数来训练模型

--batch_size 每个补丁要处理的立方体数量
--epochs 训练纪元数
--h_dim 隐藏密集层的维度（连接到变分层）
--z_dim 潜在空间维度

main.py 脚本调用模型并根据 3D CFD 数据对其进行训练。使用 NVIDIA Tesla V100 GPU 训练 100 个 epoch 大约需要 2 小时。在本例中，模型训练了 170 个 epoch。

请注意，在训练 3DConvs 模型时，与 2DConvs 模型相比，学习参数的数量呈指数级增长，因此，3D 数据的训练时间要长得多。

6、模型输出

训练完pytorch模型后，会生成一个包含训练后的权重的文件checkpoint.pkl。

在训练过程中，每隔 n 个时期根据测试数据对模型进行评估，脚本将重建的立方体与原始立方体进行比较，并将它们保存为图像。此外，损失值被记录并放置在运行文件夹中，可以通过输入以下内容使用张量板可视化损失曲线：

cd /path-to-repo/Convolutional-VAE-for-3D-Turbulence-Data
tensorboard --logdir=runs/

如果没有创建，文件夹runs/会自动生成。

7、3D 重建结果

下图中，我们展示了同一立方体样本每 n 个 epoch 的重建结果。

顶行包含原始立方体样本（对于每个速度通道）。底行包含每 n 个时期的重建输出。

对于这个例子，我们展示了从 0 到 355 个时期的重建，间隔为 15 个时期。请注意作为历元函数的重建的改进。

原文链接：Training a Convolutional Variational Autoencoder on 3D CFD Turbulence Data

BimAnt翻译整理，转载请标明出处