This project aims at providing the necessary building blocks for easily creating an image classification model using PyTorch.

Note: I finished this project in my spare time within a week. So there is still a lot of work to be done.

• Convenient: You can use a config file to create an image classification model and train on your own datasets without writing any code.
• Extensible: You can write your own modules (Dataset, Transform, Network, Loss and so on) and register them to the default config easily.
• Parameter-is-Module: You can create a module by the parameter which is consisted of a module name and an argument list.
• Multi-GPU training and inference: You can train your model on one GPU or use multi-GPU to train the model in parallel.

The top-1 accuracy (%) of different models in CIFAR-10 are shown as below. NETWORK_STRIDE is set to (2,2,2,2,2) and (1,1,2,2,2) respectively. Refer to DETAILS.md for more details about the parameter NETWORK_STRIDE.

pip3 install -r requirements.txt

We provide a helper class to simplify writing inference pipelines using pre-trained models. Here is how we would do it. Run following code from the demo folder. (The pre-trained model and sample images can be downloaded from here [Baidu(PWD: f25u)][OneDrive]. You can choose any config in ./configs)

import sys

sys.path.append("../")

import cv2

from predictor import ClsDemo
from pytorch_classification.config import cfg


config = "/path_to_config"
img_path = "/path_to_image"
checkpoint_path = "/path_to_pre-trained_model"

cfg.merge_from_file(config)
cfg.merge_from_list(["CHECKPOINT", checkpoint_path])

cls_demo = ClsDemo(cfg)

image = cv2.imread(img_path)
pred = cls_demo.run_on_openv_image(image)

You need to download the CIFAR-10 dataset and convert it the required GeneralDataset format in this codebase. (You can also download from here [Baidu(PWD: f25u)][OneDrive], which has been reformatted.) We recommend to symlink the path to the cifar-10 dataset to ./datasets as follows

# symlink the cifar-10 dataset
cd pytorch-classification
mkdir -p datasets/cifar-10
ln -s /path_to_cifar-10_dataset datasets/cifar-10

You can also configure your own paths to the datasets. For that, all you need to do is to modify ./pytorch_classification/config/data_catalog.py to point to the location where your dataset is stored.(See DETAILS.md for more details.)

You can run the following without modifications to train your model on a single GPU.

python3 tools/train.py --config-file "configs/config_cifar10_R50_1gpu.yaml"

We use internally torch.distributed.launch in order to launch multi-GPU training. This utility function from PyTorch spawns as many Python processes as the number of GPUs we want to use, and each Python process will only use a single GPU.

export NGPUS=8
python3 -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train.py --config-file "configs/config_cifar10_R50_8gpu.yaml"

If you want to train your model on more GPUs, you should change the batch size SOLVER.BATCH_SIZE and learning rate SOLVER.BASE_LR adaptively.

You can test your model directly on single or multiple GPUs. Here is an example for multi-GPU testing:

export NGPUS=8
python3 -m torch.distributed.launch --nproc_per_node=$NGPUS tools/test.py --config-file "configs/config_cifar10_R50_8gpu.yaml"

You can refer to DETAILS.md for more details.

This project is released under the MIT license. See LICENSE for additional details.

This codebase is heavily influenced by the project maskrcnn-benchmark.