Densely Connected Convolutional Networks

This is a PyTorch implementation of the DenseNet architecture as described in Densely Connected Convolutional Networks by G. Huang, Z. Liu, K. Weinberger, and L. van der Maaten.

To-do

• [ ] Multi-GPU support
• [ ] Unique model checkpointing (clashes can currently occur)

Requirements

• Python 3
• PyTorch (newest version)
• tqdm
• tensorboard_logger

Usage

This implementation currently supports training on the CIFAR-10 and CIFAR-100 datasets (support for ImageNet coming soon).

Basically, when training a model, you should specify whether to use the bottleneck variant of the dense block or not --bottleneck and if so, what compression factor to use --compression. You should also specify the total number of layers num_layers_total in the model. By default, data augmentation is performed which means no dropout, so if you choose to turn that off, you should specify a desired dropout rate --dropout_rate.

Furthermore, checkpoints of the model are saved at the end of every epoch. This means that you can resume training from your latest epoch by using the --resume=True argument. Note that this will work only after you've run at least 1 epoch of training. When testing a model, reuse whatever command you used to train the model and add the --is_train=False argument. This will load the model with the best validation accuracy and test it on the test set.

Note that you can use tensorboard to view losses and accuracy by setting the use_tensorboard argument in which case you need to run tensorboard --logdir=./logs/ in a separate shell.

Finally, to see all possible options, run:

python main.py --help

which will print:

usage: main.py [-h] [--num_blocks NUM_BLOCKS]
               [--num_layers_total NUM_LAYERS_TOTAL]
               [--growth_rate GROWTH_RATE] [--bottleneck BOTTLENECK]
               [--compression COMPRESSION] [--dataset DATASET]
               [--valid_size VALID_SIZE] [--batch_size BATCH_SIZE]
               [--num_worker NUM_WORKER] [--augment AUGMENT]
               [--shuffle SHUFFLE] [--show_sample SHOW_SAMPLE]
               [--is_train IS_TRAIN] [--epochs EPOCHS] [--init_lr INIT_LR]
               [--momentum MOMENTUM] [--weight_decay WEIGHT_DECAY]
               [--lr_decay LR_DECAY] [--dropout_rate DROPOUT_RATE]
               [--random_seed RANDOM_SEED] [--data_dir DATA_DIR]
               [--ckpt_dir CKPT_DIR] [--logs_dir LOGS_DIR] [--num_gpu NUM_GPU]
               [--use_tensorboard USE_TENSORBOARD] [--resume RESUME]
               [--print_freq PRINT_FREQ]

DenseNet

optional arguments:
  -h, --help            show this help message and exit

Network:
  --num_blocks NUM_BLOCKS
                        # of Dense blocks to use in the network
  --num_layers_total NUM_LAYERS_TOTAL
                        Total # of layers in the network
  --growth_rate GROWTH_RATE
                        Growth rate (k) of the network
  --bottleneck BOTTLENECK
                        Whether to use bottleneck layers
  --compression COMPRESSION
                        Compression factor theta in the range [0, 1]

Data:
  --dataset DATASET     Which dataset to work with. Can be CIFAR10, CIFAR100
                        or Imagenet
  --valid_size VALID_SIZE
                        Proportion of training set used for validation
  --batch_size BATCH_SIZE
                        # of images in each batch of data
  --num_worker NUM_WORKER
                        # of subprocesses to use for data loading
  --augment AUGMENT     Whether to apply data augmentation or not
  --shuffle SHUFFLE     Whether to shuffle the dataset after every epoch
  --show_sample SHOW_SAMPLE
                        Whether to visualize a sample grid of the data

Training:
  --is_train IS_TRAIN   Whether to train or test the model
  --epochs EPOCHS       # of epochs to train for
  --init_lr INIT_LR     Initial learning rate value
  --momentum MOMENTUM   Nesterov momentum value
  --weight_decay WEIGHT_DECAY
                        weight decay penalty
  --lr_decay LR_DECAY, --list LR_DECAY
                        List containing fractions of the total number of
                        epochs in which the learning rate is decayed. Enter
                        empty string if you want a constant lr.
  --dropout_rate DROPOUT_RATE
                        Dropout rate used with non-augmented datasets

Misc:
  --random_seed RANDOM_SEED
                        Seed to ensure reproducibility
  --data_dir DATA_DIR   Directory in which data is stored
  --ckpt_dir CKPT_DIR   Directory in which to save model checkpoints
  --logs_dir LOGS_DIR   Directory in which Tensorboard logs wil be stored
  --num_gpu NUM_GPU     # of GPU's to use. A value of 0 will run on the CPU
  --use_tensorboard USE_TENSORBOARD
                        Whether to use tensorboard for visualization
  --resume RESUME       Whether to resume training from most recent checkpoint
  --print_freq PRINT_FREQ
                        How frequently to display training details on screen

You can edit the default values of these arguments in the config.py file.

Here's an example command for training a DenseNet-BC-100 architecture with a growth rate of 12, data augmentation, tensorboard visualization and with GPU:

python main.py \
--num_layers_total=100 \
--bottleneck=True \
--compression=0.5 \
--num_gpu=1 \
--use_tensorboard=True

Performance

I trained DenseNet-40 and DenseNet-BC-100 variants on the CIFAR-10 dataset but was not able to reproduce the author's results. I don't know if this stems from an error in the implementation or an unlucky seed... Training the 2 models in parallel took 2 days to complete on a p2.xlarge AWS instance with 1 GPU. If I get the time, I'll create clean, minimal instructions for setting up a similar instance for free.

Here are some tensorboard visualizations comparing the two models:

From looking at the losses and accuracies, it's clearly visible that decreasing the learning rate at earlier times than mentioned in the paper can shorten the training time by a large factor. In fact, I noticed during training that the train accuracy and loss would just stagnate for a dozen epochs and then have a significant jump when the learning rate was decreased halway through. I'll be testing out this intuition and report my findings at a later date.

References

• Thanks to Taehoon Kim for inspiring the general file hierarchy and layout of this project.
• Thanks to the PyTorch ImageNet training example for helping me code the Trainer class.