A Simple Pipeline to Train PyTorch FasterRCNN Model

Train PyTorch FasterRCNN models easily on any custom dataset. Choose between official PyTorch models trained on COCO dataset, or choose any backbone from Torchvision classification models, or even write your own custom backbones.

You can run a Faster RCNN model with Mini Darknet backbone and Mini Detection Head at more than 150 FPS on an RTX 3080.

Get Started

​

• Find blog posts/tutorials on DebuggerCafe

Updates

• June 6 2025: Support for both Pascal VOC and YOLO text file annotation type during training. Check custom training section

• August 28 2024: SAHI image inference for all pretrained Torchvision Faster RCNN models integrated. Find the script here.

• Filter classes to visualize during inference using the --classes command line argument with space separated class indices from the dataset YAML file.

  For example, to visualize only persons in COCO dataset, use, python inference.py --classes 1 <rest of the command>

  To visualize person and car, use, python inference.py --classes 1 3 <rest of the command>

• Added Deep SORT Real-Time tracking to inference_video.py and onnx_video_inference.py. Using --track command with the usual inference command. Support for MobileNet Re-ID for now.

Custom Model Naming Conventions

For this repository:

• Small head refers to 512 representation size in the Faster RCNN head and predictor.
• Tiny head refers to 256 representation size in the Faster RCNN head and predictor.
• Nano head refers to 128 representation size in the Faster RCNN head and predictor.

Check All Available Model Flags

Go To

• Setup on Ubuntu
• Setup on Windows
• Train on Custom Dataset
• Inference
• Evaluation
• Available Models
• Tutorials

Setup on Ubuntu

1. Clone the repository.

   git clone https://github.com/sovit-123/fastercnn-pytorch-training-pipeline.git
   

   Optional: Initialize DINOv3 submodule for training DINOv3 Faster RCNN models.

   git submodule update --init
   

2. Install requirements as per GPU. Install requirements on RTX 30/40 (Ampere and Ada Lovelace) series and T4/P100 GPUs.

   pip install -r requirements.txt
   

OR

Install requirements for RTX 50 series and Blackwell GPUs. First install PyTorch >= 2.8 with CUDA >= 12.9

pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu129

Install rest of the requirements

pip install -r requirements_blackwell.txt

Setup on Windows

1. First you need to install Microsoft Visual Studio from here. Sing In/Sing Up by clicking on this link and download the Visual Studio Community 2017 edition.

   

   Install with all the default chosen settings. It should be around 6 GB. Mainly, we need the C++ Build Tools.

2. Then install the proper pycocotools for Windows.

   pip install git+https://github.com/gautamchitnis/cocoapi.git@cocodataset-master#subdirectory=PythonAPI
   

3. Clone the repository.

   git clone https://github.com/sovit-123/fastercnn-pytorch-training-pipeline.git
   

4. Then install the remaining requirements except for pycocotools.

   Install requirements on RTX 30/40 (Ampere and Ada Lovelace) series and T4/P100 GPUs.

   pip install -r requirements.txt
   

OR

Install requirements for RTX 50 series and Blackwell GPUs. First install PyTorch >= 2.8 with CUDA >= 12.9

pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu129

Install rest of the requirements (apart from pycocotools)

pip install -r requirements_blackwell.txt

Using Custom Weights

Some models like DINOv3 based Faster RCNN models require the pretrained weights to be present locally. Put all the DINOv3 backbone weights in the weights directory. The respective model files will load them from the directory.

You can download the weights by filling the form here.

For example, for the FasterRCNN DINOv3 ConvNext Tiny model (models/fasterrcnn_dinov3_convnext_tiny.py) the weights are loaded using the following syntax with relative path.

# Relative to parent fasterrcnn directory.
REPO_DIR = 'dinov3'
# Relative to parent fasterrcnn directory or the absolute path.
WEIGHTS_URL = 'weights/dinov3_convnext_tiny_pretrain_lvd1689m-21b726bb.pth'

self.backbone = torch.hub.load(
    REPO_DIR, 
    "dinov3_convnext_tiny", 
    source='local', 
    weights=WEIGHTS_URL
)

Train on Custom Dataset

Taking an exmaple of the smoke dataset from Kaggle. Let's say that the dataset is in the data/smoke_pascal_voc directory in the following format. And the smoke.yaml is in the data_configs directory. Assuming, we store the smoke data in the data directory

├── data
│   ├── smoke_pascal_voc
│   │   ├── archive
│   │   │   ├── train
│   │   │   └── valid
│   └── README.md
├── data_configs
│   └── smoke.yaml
├── models
│   ├── create_fasterrcnn_model.py
│   ...
│   └── __init__.py
├── outputs
│   ├── inference
│   └── training
│       ...
├── readme_images
│   ...
├── torch_utils
│   ├── coco_eval.py
│   ...
├── utils
│   ├── annotations.py
│   ...
├── datasets.py
├── inference.py
├── inference_video.py
├── __init__.py
├── README.md
├── requirements.txt
└── train.py

The content of the smoke.yaml should be the following. The folder containing the annotation files can either point to Pascal VOC XML files or YOLO text labels folder. The images and labels (for both Pascal VOC XML and YOLO text files) can be either in the same folder or in different folders because the image and annotation files are matched based on the file names during dataset preparation.

If the data config file (shown below) points to Pascal VOC XML annotations, the CLASSES field can contain the class names in any order after the __background__ class. If the data config file points to YOLO text file annotation folder, the CLASSES should contain the class names in the order as present in the YOLO dataset data.yaml file. This is necessary to maintain indexing order during training.

# Images and labels direcotry should be relative to train.py
TRAIN_DIR_IMAGES: ../../xml_od_data/smoke_pascal_voc/archive/train/images
TRAIN_DIR_LABELS: ../../xml_od_data/smoke_pascal_voc/archive/train/annotations # This can contain .xml or .txt files
# VALID_DIR should be relative to train.py
VALID_DIR_IMAGES: ../../xml_od_data/smoke_pascal_voc/archive/valid/images
VALID_DIR_LABELS: ../../xml_od_data/smoke_pascal_voc/archive/valid/annotations # This can contain .xml or .txt files

# Class names.
CLASSES: [
    '__background__',
    'smoke'
]

# Number of classes (object classes + 1 for background class in Faster RCNN).
NC: 2

# Whether to save the predictions of the validation set while training.
SAVE_VALID_PREDICTION_IMAGES: True

Note that the data and annotations can be in the same directory as well. In that case, the TRAIN_DIR_IMAGES and TRAIN_DIR_LABELS will save the same path. Similarly for VALID images and labels. The datasets.py will take care of that.

Next, to start the training, you can use the following command.

Command format:

During training, we need to provide a --label-type argument which should be either yolo or pascal_voc depending on the annotation folder path in the data configuration file above. Default is pascal_voc

python train.py --data <path to the data config YAML file> --epochs 100 --model <model name (defaults to fasterrcnn_resnet50)> --name <folder name inside output/training/> --batch 16 --label-type <pascal_voc or yolo>

In this case, the exact command would be:

python train.py --data data_configs/smoke.yaml --epochs 100 --model fasterrcnn_resnet50_fpn --name smoke_training --batch 16 --label-type pascal_voc

The terimal output should be similar to the following:

Number of training samples: 665
Number of validation samples: 72

3,191,405 total parameters.
3,191,405 training parameters.
Epoch     0: adjusting learning rate of group 0 to 1.0000e-03.
Epoch: [0]  [ 0/84]  eta: 0:02:17  lr: 0.000013  loss: 1.6518 (1.6518)  time: 1.6422  data: 0.2176  max mem: 1525
Epoch: [0]  [83/84]  eta: 0:00:00  lr: 0.001000  loss: 1.6540 (1.8020)  t