add data download files; minor cleanup

.gitignore

@@ -1,3 +1,6 @@
 # Python
 *.pyc
 *.swp
+
+# data files
+*.h5

README.md

@@ -12,17 +12,20 @@ This repository contains source code for the experiments in [LCA: Loss Change Al
 }
 ```
 
+To read more about this project, see the [Uber Eng blog post](https://eng.uber.com/loss-change-allocation).
+
 ## Codebase structure
 
+* `data/download_mnist.py` and `data/download_cifar10.py` downloads MNIST/CIFAR10 data and splits it into train, val, and test, and saves them in the data folder as h5 files
 * `train.py` trains a model and saves weights at every iteration, using architectures defined in `network_builders.py`
-* `adaptive_calc_gradients.py` calculates gradients used for LC, based on saved weights
-* `save_lc_stream.py` calculates LC for runs where gradients or weights don't fit into memory
+* `adaptive_calc_gradients.py` calculates gradients used for LCA, based on saved weights
+* `save_lca_stream.py` calculates LCA for runs where gradients or weights don't fit into memory
 * `plot_util.py` demonstrates some examples of various visualizations, aggregations, and analyses
 
 (Optional) This codebase uses the `GitResultsManager` package to keep track of experiments. See: https://github.com/yosinski/GitResultsManager
 
 
 ## Example commands
 
-* To train: `python train.py --train_h5 /data/cifar10_train --test_h5 /data/cifar10_test --input_dim 32,32,3 --arch resnet --opt sgd --lr 0.1 --save_weights --num_epochs 25 --large_batch_size 5000 --test_batch_size 5000 --eval_every 100 --print_every 100 --log_every 50`
-* To calculate gradients: `python adaptive_calc_gradients.py --train_h5 /data/cifar10_train --test_h5 /data/cifar10_test --input_dim 32,32,3 --arch resnet --opt sgd --large_batch_size 2500 --test_batch_size 2500 --num_gpus 4 --weights_h5 previous_experiment/weights`
+* To train: `python train.py --train_h5 data/cifar10_train.h5 --test_h5 data/cifar10_val.h5 --input_dim 32,32,3 --arch resnet --opt sgd --lr 0.1 --save_weights --num_epochs 25 --large_batch_size 5000 --test_batch_size 5000 --eval_every 100 --print_every 100 --log_every 50 --output_dir outputs`
+* To calculate gradients: `python adaptive_calc_gradients.py --train_h5 data/cifar10_train.h5 --test_h5 data/cifar10_val.h5 --input_dim 32,32,3 --arch resnet --opt sgd --large_batch_size 2500 --test_batch_size 2500 --num_gpus 4 --weights_h5 outputs/weights`

data/download_cifar10.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2019 Uber Technologies, Inc.
+
+# Licensed under the Uber Non-Commercial License (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at the root directory of this project.
+
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tensorflow as tf
+import h5py
+import numpy as np
+
+np.random.seed(seed=0)
+
+def main():
+    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
+
+    valset_ind = np.random.choice(range(50000), size=5000, replace=False)
+    trainset_ind = np.array([i for i in range(50000) if i not in valset_ind])
+
+    train_set_images = x_train[trainset_ind]
+    train_set_images = train_set_images.reshape((45000,32,32,3))
+    train_set_labels = y_train[trainset_ind]
+    train_set_labels = train_set_labels.reshape(train_set_labels.shape[0])
+
+    val_set_images = x_train[valset_ind]
+    val_set_images = val_set_images.reshape((5000,32,32,3))
+    val_set_labels = y_train[valset_ind]
+    val_set_labels = val_set_labels.reshape(val_set_labels.shape[0])
+
+    x_test = x_test.reshape((10000,32,32,3))
+    y_test = y_test.reshape(y_test.shape[0])
+
+    f = h5py.File("cifar10_train.h5", "w")
+    f.create_dataset('images', data=train_set_images)
+    f.create_dataset('labels', data=train_set_labels)
+    f.close()
+
+    f = h5py.File("cifar10_val.h5", "w")
+    f.create_dataset('images', data=val_set_images)
+    f.create_dataset('labels', data=val_set_labels)
+    f.close()
+
+    f = h5py.File("cifar10_test.h5", "w")
+    f.create_dataset('images', data=x_test)
+    f.create_dataset('labels', data=y_test)
+    f.close()
+
+if __name__ == '__main__':
+    main()

data/download_mnist.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2019 Uber Technologies, Inc.
+
+# Licensed under the Uber Non-Commercial License (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at the root directory of this project. 
+
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tensorflow as tf
+import h5py
+import numpy as np
+
+np.random.seed(seed=0)
+
+def main():
+    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
+
+    valset_ind = np.random.choice(range(60000), size=5000, replace=False)
+    trainset_ind = np.array([i for i in range(60000) if i not in valset_ind])
+
+    train_set_images = x_train[trainset_ind]
+    train_set_images = train_set_images.reshape((55000,28,28,1))
+    train_set_labels = y_train[trainset_ind]
+
+    val_set_images = x_train[valset_ind]
+    val_set_images = val_set_images.reshape((5000,28,28,1))
+    val_set_labels = y_train[valset_ind]
+
+    x_test = x_test.reshape((10000,28,28,1))
+
+    f = h5py.File("mnist_train.h5", "w")
+    f.create_dataset('images', data=train_set_images)
+    f.create_dataset('labels', data=train_set_labels)
+    f.close()
+
+    f = h5py.File("mnist_val.h5", "w")
+    f.create_dataset('images', data=val_set_images)
+    f.create_dataset('labels', data=val_set_labels)
+    f.close()
+
+    f = h5py.File("mnist_test.h5", "w")
+    f.create_dataset('images', data=x_test)
+    f.create_dataset('labels', data=y_test)
+    f.close()
+
+if __name__ == '__main__':
+    main()

plot_util.py

@@ -28,8 +28,6 @@
 from scipy.ndimage import gaussian_filter1d
 from scipy.stats import dweibull
 from ast import literal_eval
-from IPython import embed
-import pdb
 
 ######################################## handling large matrices ########################################
 

save_lc_stream.py → save_lca_stream.py

@@ -20,12 +20,7 @@
 from ast import literal_eval
 import sys
 
-lab_root = os.path.join(os.path.abspath(os.path.dirname(__file__)), '..', '..')
-sys.path.insert(1, lab_root)
-
 from general.util import mkdir_p
-import matplotlib
-matplotlib.use('Agg')
 import plot_util as util
 
 def make_parser():