udpate

README.md

@@ -4,13 +4,13 @@
 This repository contains the PyTorch code for the paper
 > Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer and Tom Goldstein. [*Visualizing the Loss Landscape of Neural Nets*](https://arxiv.org/abs/1712.09913). NIPS, 2018.
 
-Given a network architecture and its pre-trained model, this tool calculates and visualizes the model's surrounding loss surface along random direction(s) on the training set.
+Given a network architecture and its pre-trained parameters, this tool calculates and visualizes the model's surrounding loss surface along random direction(s) on the training set.
 The calculation can be done in parallel with multiple GPUs with multiple nodes.
-The direction(s) and the surface values are saved in HDF5 (`.h5`) files.
+The direction(s) and the surface values are stored in HDF5 (`.h5`) files.
 
 ## Setup
 
-**Environment**: A (multi-) GPU node with following software/libraries installed:
+**Environment**: One multi-GPU node with following software/libraries installed:
 - [PyTorch 0.4](https://pytorch.org/)
 - [openmpi 3.1.2](https://www.open-mpi.org/)
 - [mpi4py 2.0.0](https://mpi4py.scipy.org/docs/usrman/install.html)
@@ -21,13 +21,11 @@ The direction(s) and the surface values are saved in HDF5 (`.h5`) files.
 **Pre-trained models**:
 The code accepts pre-trained PyTorch models for CIFAR-10 dataset.
 To load the pre-trained model correctly, the model file should contain `state_dict`, which is saved from the `state_dict()` method.
-The default saving folder for pre-trained networks is `cifar10/trained_nets/`.
-
-The pre-trained models can be downloaded [here](https://www.cs.umd.edu/~tomg/projects/landscapes/):
-
+The default path for pre-trained networks is `cifar10/trained_nets/`.
+Some of the pre-trained models and plotted figures can be downloaded [here](http://www.cs.umd.edu/~haoli/projects/landscape/trained_net.tar.gz) (464 MB).
 
 **Data preprocessing**:
-The data normalization for visualization should be consistent with the one used for model training.
+The data normalization method for visualization should be consistent with the one used for model training.
 No data augmentation (random cropping or horizontal flipping) is used in calculating the loss values.
 
 ## Visualizing 1D loss curve
@@ -61,7 +59,7 @@ mpirun -n 4 python plot_surface.py --mpi --cuda --model vgg9 --x=-1:1:51 \
 
  ![VGG-9 SGD, WD=0](doc/images/vgg9_sgd_lr=0.1_bs=128_wd=0_300.t7_states_ignore=biasbn_norm=filter_[-1.0,1.0,51].h5_1d_loss_err.jpg)
 
-One can also customize the 1D plots with `plot_1D.py` once the suface file is available.
+We can also customize the 1D plots with `plot_1D.py` once the surface file is available.
 
 
 ## Visualizing 2D loss contours
@@ -76,10 +74,10 @@ mpirun -n 4 python plot_surface.py --model resnet56 --x=-1:1:51 --y=-1:1:51 \
 
 ![ResNet-56](doc/images/resnet56_random_[-1.0,1.0]x[-1.0,1.0].h5_2dcontour.jpg)
 
-We can also customize the plots given a surface `.h5` file with `plot_2D.py`, which supports plotting the loss/accuracy/error surface for both training and validation set.
+We can also customize the plots given a surface `.h5` file with `plot_2D.py`.
 
 ```
-python plot_2D.py --file path_to_surface_file --surface_name train_loss
+python plot_2D.py --surf_file path_to_surface_file --surface_name train_loss
 ```
 - `--surface_name` specifies the type of surface. The default choice is `train_loss`,
 - `--vmin` and `--vmax` sets the range of values to be plotted.
@@ -93,20 +91,17 @@ You may also want to render the 3D surface with [ParaView](http://paraview.org).
 ![ResNet-56-noshort](doc/images/resnet56_noshort_small.jpg)
 ![ResNet-56](doc/images/resnet56_small.jpg)
 
-1. Convert surface `.h5` file to `.vtp` file.
+1. Convert the surface `.h5` file to the `.vtp` file.
 ```
-python h52vtp.py --file path_to_surf_file --zmax  10 --log
+python h52vtp.py --file path_to_surface_file --surf_name train_loss --zmax  10 --log
 ```
- - It will generate a [VTK](https://www.kitware.com/products/books/VTKUsersGuide.pdf) file containing the loss surface with max value 10 in the log scale.
+which will generate a [VTK](https://www.kitware.com/products/books/VTKUsersGuide.pdf) file containing the loss surface with max value 10 in the log scale.
 
-2. Open the `.vtp` file with ParaView
- - In ParaView, open the `.vtp` file with the VTK reader. Click the eye icon in the `Pipeline Browser` to make the figure show up.
- - You can drag the surface around, and change the colors in the `Properties` window.  
- - `Save screenshot` in the File menu saves the image, which can be cropped elsewhere.
+2. Open the `.vtp` file with ParaView. In ParaView, open the `.vtp` file with the VTK reader. Click the eye icon in the `Pipeline Browser` to make the figure show up. You can drag the surface around, and change the colors in the `Properties` window.  `Save screenshot` in the File menu saves the image, which can be cropped elsewhere.
 
 
 ## Citation
-If you find this repo useful in your research, please cite:
+If you find this code useful in your research, please cite:
 
 ```
 @inproceedings{visualloss,