神经网络损失景观可视化

本存储库包含了论文中的 PyTorch 代码：

Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer 和 Tom Goldstein. 《Visualizing the Loss Landscape of Neural Nets》. NIPS, 2018.

telesens 提供了一个交互式的 3D 损失表面可视化器：http://www.telesens.co/loss-landscape-viz/viewer.html

给定一个网络架构及其预训练参数，此工具可以计算并沿着最优参数附近的随机方向（s）可视化损失表面。计算可以在每个节点上的多个 GPU 并行进行，并且支持多节点。产生的随机方向（s）和损失表面值被存储在 HDF5（.h5）文件中。

设置

环境：至少有一个安装了以下软件/库的多 GPU 节点（s）：

• PyTorch 0.4
• openmpi 3.1.2
• mpi4py 2.0.0
• numpy 1.15.1
• h5py 2.7.0
• matplotlib 2.0.2
• scipy 0.19

预训练模型： 代码接受针对 CIFAR-10 数据集的预训练 PyTorch 模型。为了正确加载预训练模型，模型文件应包含从 state_dict() 方法保存的 state_dict。预训练网络的默认路径是 cifar10/trained_nets。一些预训练模型和绘制的图形可以从这里下载：

• VGG-9 (349 MB)
• ResNet-56 (10 MB)
• ResNet-56-noshort (20 MB)
• DenseNet-121 (75 MB)

数据预处理： 用于可视化的数据预处理方法应与模型训练时的方法一致。在计算损失值时不使用数据增强（如随机裁剪或水平翻转）。

可视化 1D 损失曲线

创建 1D 线性插值

1D 线性插值方法 [1] 在同一网络损失函数两个最小值之间的方向上评估损失值。这种方法已用于比较不同批量大小下训练出的最小值的平坦度 [2]。通过使用 plot_surface.py 方法生成 1D 线性插值图。

mpirun -n 4 python plot_surface.py --mpi --cuda --model vgg9 --x=-0.5:1.5:401 --dir_type states \
--model_file cifar10/trained_nets/vgg9_sgd_lr=0.1_bs=128_wd=0.0_save_epoch=1/model_300.t7 \
--model_file2 cifar10/trained_nets/vgg9_sgd_lr=0.1_bs=8192_wd=0.0_save_epoch=1/model_300.t7 --plot

• --x=-0.5:1.5:401 设置绘图范围和分辨率。绘图中的 x 坐标将从 -0.5 到 1.5（最小值位于 0 和 1），沿这条线评估损失值的 401 个位置。
• --dir_type states 表示方向包含所有参数以及 BN 层的统计信息（running_mean 和 running_var）。请注意，在同一个图中一起绘制两个解决方案时，忽略 running_mean 和 running_var 无法产生正确的损失值。
• 两个模型文件包含描述损失函数两个不同最小值的网络参数。该图将在这两个极小值之间进行插值。

沿随机归一化方向生成图像

创建与模型参数具有相同维度的随机方向，并对其进行“滤波归一化”。然后我们可以在该方向上采样损失值。

mpirun -n 4 python plot_surface.py --mpi --cuda --model vgg9 --x=-1:1:51 \
--model_file cifar10/trained_nets/vgg9_sgd_lr=0.1_bs=128_wd=0.0_save_epoch=1/model_300.t7 \
--dir_type weights --xnorm filter --xignore biasbn --plot

• --dir_type weights 表明方向与学习到的参数具有相同的维度，包括偏置和 BN 层的参数。
• --xnorm filter 在过滤器级别对随机方向进行归一化。这里的“过滤器”指的是产生单个特征映射的参数。对于全连接层，“过滤器”包含贡献给单个神经元的权重。
• --xignore biasbn 忽略与偏置和 BN 参数对应的方向（在随机向量的相应条目中填充零）。

一旦表面文件可用，我们还可以通过调用 plot_1D.py 自定义 1D 图的外观。

可视化 2D 损失轮廓

要绘制损失轮廓，我们选择两个随机方向，并以与 1D 绘图相同的方式对它们进行归一化。

mpirun -n 4 python plot_surface.py --mpi --cuda --model resnet56 --x=-1:1:51 --y=-1:1:51 \
--model_file cifar10/trained_nets/resnet56_sgd_lr=0.1_bs=128_wd=0.0005/model_300.t7 \
--dir_type weights --xnorm filter --xignore biasbn --ynorm filter --yignore biasbn --plot

一旦生成表面并存储为 .h5 文件，我们可以使用脚本 plot_2D.py 生成和自定义轮廓图。

python plot_2D.py --surf_file path_to_surf_file --surf_name train_loss

• --surf_name 指定表面类型。默认选项是 train_loss，
• --vmin 和 --vmax 设置要绘制的值的范围。
• --vlevel 设置轮廓的步长。

可视化3D损失曲面

使用plot_2D.py脚本可借助matplotlib绘制基础的3D损失曲面图。 如果希望呈现更详细的渲染效果，利用光照显示细节，你可以借助ParaView来绘制损失曲面。

实现这一目标需要以下步骤：

1. 将表面.h5文件转换为.vtp文件。

python h52vtp.py --surf_file 路径_to_surf_file --surf_name train_loss --zmax 10 --log

这将生成包含对数尺度下最大值为10的损失表面的VTK文件。

2. 使用ParaView打开.vtp文件。在ParaView中，使用VTK读取器打开文件。在“Pipeline Browser”中点击眼睛图标使图像显示出来。你可以拖动表面，并在“Properties”窗口中更改颜色。

3. 如果表面显得非常瘦长像针状，你可能需要调整左侧控制面板中的“变换”参数。在“比例”字段中输入大于1的数字以加宽图表。

4. 从文件菜单中选择“保存屏幕截图”来保存图像。

参考文献

[1] Ian J Goodfellow, Oriol Vinyals, and Andrew M Saxe. Qualitatively characterizing neural network optimization problems. ICLR, 2015.

[2] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. ICLR, 2017.

引用

如果你的研究中使用了这个代码，请引用：

@inproceedings{visualloss,
  title={Visualizing the Loss Landscape of Neural Nets},
  author={Li, Hao and Xu, Zheng and Taylor, Gavin and Studer, Christoph and Goldstein, Tom},
  booktitle={Neural Information Processing Systems},
  year={2018}
}