ImageNet

This implements training of popular model architectures, such as AlexNet, SqueezeNet, ResNet, DenseNet and VGG on the ImageNet dataset(Now we supported alexnet, vgg, resnet, squeezenet, densenet).

Requirements

• PyTorch 0.4.0
• cuda && cudnn
• Download the ImageNet dataset and move validation images to labeled subfolders
  • To do this, you can use the following script: https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh

Training

To train a model, run main.py with the desired model architecture and the path to the ImageNet dataset:

python main.py [imagenet-folder with train and val folders] -a alexnet --lr 0.01

The default learning rate schedule starts at 0.01 and decays by a factor of 10 every 30 epochs.

Usage

usage: main.py [-h] [-a ARCH] [--epochs N] [--start-epoch N] [-b N] [--lr LR]
               [--momentum M] [--weight-decay W] [-j N] [-m] [-p]
               [--print-freq N] [--resume PATH] [-e]
               DIR

PyTorch ImageNet Training

positional arguments:
  DIR                   path to dataset

optional arguments:
  -h, --help            show this help message and exit
  -a ARCH, --arch ARCH  model architecture: alexnet | squeezenet1_0 |
                        squeezenet1_1 | densenet121 | densenet169 |
                        densenet201 | densenet201 | densenet161 | vgg11 |
                        vgg11_bn | vgg13 | vgg13_bn | vgg16 | vgg16_bn | vgg19
                        | vgg19_bn | resnet18 | resnet34 | resnet50 |
                        resnet101 | resnet152 (default: alexnet)
  --epochs N            numer of total epochs to run
  --start-epoch N       manual epoch number (useful to restarts)
  -b N, --batch-size N  mini-batch size (default: 256)
  --lr LR, --learning-rate LR
                        initial learning rate
  --momentum M          momentum
  --weight-decay W, --wd W
                        Weight decay (default: 1e-4)
  -j N, --workers N     number of data loading workers (default: 4)
  -m, --pin-memory      use pin memory
  -p, --pretrained      use pre-trained model
  --print-freq N, -f N  print frequency (default: 10)
  --resume PATH         path to latest checkpoitn, (default: None)
  -e, --evaluate        evaluate model on validation set

Result

The results of a single model on ILSVRC-2012 validation set.

Acknowledgement

PyTorch Examples

AlexNet

VGG

ResNet

SqueezeNet

DenseNet