Update README.md

Author: gpleiss

README.md

@@ -1,9 +1,10 @@
 # efficient_densenet_pytorch
-A PyTorch implementation of DenseNets, optimized to save GPU memory.
+A PyTorch 0.4 implementation of DenseNets, optimized to save GPU memory.
 
 ## Recent updates
-1. **Now works on PyTorch 0.3.x!**
-1. `models/densenet_efficient_multi_gpu.py` is now depricated. `models/densenet_efficient.py` can handle both single and multi-GPU setups.
+1. **Now works on PyTorch 0.4!** It uses the checkpointing feature, which makes this code WAY more efficient!!!
+1. **No longer works with PyTorch 0.3.x.** If you are using an old version of PyTorch, check out [the pytorch_0.3.1 tag](https://github.com/gpleiss/efficient_densenet_pytorch/tree/pytorch_0.3.1) for an older version of this code.
+1. `models/densenet_efficient.py` is now depricated. `models/densenet.py` can handle both the efficient and original implementations.
 
 ## Motivation
 While DenseNets are fairly easy to implement in deep learning frameworks, most
@@ -23,21 +24,25 @@ For more details, please see the [technical report](https://arxiv.org/pdf/1707.0
 ![Diagram of implementation](https://raw.github.com/gpleiss/efficient_densenet_pytorch/master/images/forward.png)
 
 ## Requirements
-- PyTorch 0.3.x
+- PyTorch 0.4
 - CUDA
 
+**N.B.** If you are using PyTorch 0.3.x, please check out [the 0.3.x compatible version](https://github.com/gpleiss/efficient_densenet_pytorch/tree/pytorch_0.3.1).
+
 ## Usage
 
 **In your existing project:**
-There are two files in the `models` folder.
- - `models/densenet.py` is a "naive" implementation, based off the [torchvision](https://github.com/pytorch/vision/blob/master/torchvision/models/densenet.py) and
+There is one file in the `models` folder.
+ - `models/densenet.py` is an implementation based off the [torchvision](https://github.com/pytorch/vision/blob/master/torchvision/models/densenet.py) and
 [project killer](https://github.com/felixgwu/img_classification_pk_pytorch/blob/master/models/densenet.py) implementations.
- - `models/densenet_efficient.py` is the new efficient implementation. (Code is still a little ugly. We're working on cleaning it up!)
-Copy either one of those files into your project!
+
+If you care about speed, and memory is not an option, pass the `efficient=False` argument into the `DenseNet` constructor.
+Otherwise, pass in `efficient=True`.
 
 **Options:**
 - All options are described in [the docstrings of the model files](https://github.com/gpleiss/efficient_densenet_pytorch/blob/master/models/densenet_efficient.py#L189)
 - The depth is controlled by `block_config` option
+- `efficient=True` uses the memory-efficient version
 - If you want to use the model for ImageNet, set `small_inputs=False`. For CIFAR or SVHN, set `small_inputs=True`.
 
 **Running the demo:**
@@ -47,7 +52,13 @@ The only extra package you need to install is [python-fire](https://github.com/g
 pip install fire
 ```
 
-- single GPU:
+- Single GPU:
+
+```sh
+CUDA_VISIBLE_DEVICES=0 python demo.py --efficient True --data <path_to_folder_with_cifar10> --save <path_to_save_dir>
+```
+
+- Multiple GPU:
 
 ```sh
 CUDA_VISIBLE_DEVICES=0,1,2 python demo.py --efficient True --data <path_to_folder_with_cifar10> --save <path_to_save_dir>