# RevNet with TensorFlow eager execution

This folder contains a TensorFlow eager implementation of the [Reversible Residual Network](https://arxiv.org/pdf/1707.04585.pdf) adapted from the released implementation by the authors. The presented implementation can be ran with both eager and graph execution. The code is considerably simplified with `tf.GradientTape`. Moreover, we reduce the a redundant forward pass in the implementation by the authors. This saves us from using `tf.stop_gradient` and makes the model run faster.

##  Content

- `revnet.py`: The RevNet model.
- `blocks.py`: The relevant reversible blocks.
- `ops.py`: Auxiliary downsampling operation.
- `cifar_tfrecords.py`: Script to generate the TFRecords for both CIFAR-10 and CIFAR-100.
- `cifar_input.py`: Script to read from TFRecords and generate dataset objects with the `tf.data` API.
- `config.py`: Configuration file for network architectures and training hyperparameters.
- `main.py`: Main training and evaluation script.
- `main_estimator.py`: Script to train RevNet models on CIFAR-10 and CIFAR-100 with the `tf.estimator` API.
- `main_estimator_tpu.py`: Script to train RevNet models on ImageNet with TPU estimators on Cloud TPUs.
- `resnet_preprocessing.py`, `imagenet_input.py`: Boilerplate to read ImageNet data from TFRecords.

## Train on CIFAR-10/CIFAR-100
- Make sure you have installed TensorFlow 1.10+ or the latest `tf-nightly`
or `tf-nightly-gpu` pip package in order to access the eager execution feature.

- First run

```bash
python cifar_tfrecords.py --data_dir ${PWD}/cifar
```
to download the cifar dataset and convert them
to TFRecords. This produces TFRecord files for both CIFAR-10 and CIFAR-100.

- To train a model, run

```bash
python main.py --data_dir ${PWD}/cifar
```

- Optional arguments for `main.py` include
  - `train_dir`: Directory to store eventfiles and checkpoints.
  - `restore`: Restore the latest checkpoint.
  - `validate`: Use validation set for training monitoring.
  - `dataset`: Use either `cifar-10` or `cifar-100`.
  - `config`: RevNet configuration.
  - `use_defun`: Use `tfe.defun` to boost performance.

- To train a model with estimators in graph execution, run

```bash
python main_estimator.py --data_dir ${PWD}/cifar
```
To ensure our code works properly when using the Keras model in an estimator,
`tf-nightly` or `tf-nightly-gpu` is highly recommended as of August 2018.

- Optional arguments for `main.py` include
  - `model_dir`: Directory to store eventfiles and checkpoints.
  - `dataset`: Use either `cifar-10` or `cifar-100`.
  - `config`: RevNet configuration.
  - `export`: Export the model for serving if True.

## Speed up with `tfe.defun`
To ensure that `tf.contrib.eager.defun` in our code works properly with all
part of the model during training, the latest `tf-nightly` or `tf-nightly-gpu`
is highly recommended as of August 2018.

Even though the speed difference between pure eager execution and graph execution is noticeable,
the difference between fully "defunned" model training and graph
training is negligible.

## Train on ImageNet with Cloud TPUs
The standard way to train models on Cloud TPUs is via TPU estimators and graph
execution. Models built with the `tf.keras` API are fully compatible with TPU estimators.
To ensure our code works properly in this setting,
`tf-nightly` or `tf-nightly-gpu` is highly recommended as of August 2018.

### Setup a Google Cloud project

Follow the instructions at the [Quickstart Guide](https://cloud.google.com/tpu/docs/quickstart)
to get a GCE VM with access to Cloud TPU.

To run this model, you will need:

* A GCE VM instance with an associated Cloud TPU resource
* A GCS bucket to store your training checkpoints
* (Optional): The ImageNet training and validation data preprocessed into
  TFRecord format, and stored in GCS.

### Format the data

The data is expected to be formatted in TFRecord format, as generated by [this
script](https://github.com/tensorflow/tpu/blob/master/tools/datasets/imagenet_to_gcs.py).

If you do not have ImageNet dataset prepared, you can use a randomly generated
fake dataset to test the model. It is located at
`gs://cloud-tpu-test-datasets/fake_imagenet`.

### Start training

Train the model by executing the following command (substituting the appropriate
values):

```bash
python main_estimator_tpu.py \
  --tpu=$TPU_NAME \
  --data_dir=$DATA_DIR \
  --model_dir=$MODEL_DIR
```

## Performance
- RevNet-38 achieves >92% and >71% accuracy on CIFAR-10 and CIFAR-100 respectively.
- RevNet-56 achieves <26% top-1 error rate on ImageNet.

## Reference
The Reversible Residual Network: Backpropagation Without Storing Activations.
Aidan N. Gomez, Mengye Ren, Raquel Urtasun, Roger B. Grosse. Neural Information Processing Systems (NIPS), 2017.
