AL-MDN: Active Learning for Deep Object Detection

repository·main·Indexed 41 days ago

https://github.com/nvlabs/al-mdn

A PyTorch implementation of Active Learning for Deep Object Detection using Probabilistic Modeling. AL-MDN utilizes mixture density networks to estimate localization and classification uncertainties in a single forward pass as an alternative to ensembles or MC-Dropout. Built upon ssd.pytorch, it supports supervised and active learning training scripts and evaluation on PASCAL VOC and MS-COCO datasets.

Tokens
612
Snippets
4
Records
5
Agent score
2%

What's inside AL-MDN

  1. Train the model (Supervised or Active Learning)

    main

    Run the training scripts using the CUDA_VISIBLE_DEVICES environment variable to specify the GPU ID.

    • Supervised learning: Use train_ssd_gmm_supervised_learning.py.
    • Active learning: Use train_ssd_gmm_active_learining.py.
    # Supervised learning
    CUDA_VISIBLE_DEVICES=<GPU_ID> python train_ssd_gmm_supervised_learning.py
    
    # Active learning
    CUDA_VISIBLE_DEVICES=<GPU_ID> python train_ssd_gmm_active_learining.py
  2. Evaluate the trained model on PASCAL VOC or MS-COCO

    main

    To evaluate your model, use the corresponding evaluation script and provide the path to your trained weights using the --trained_model flag.

    Note for MS-COCO: You must update COCO_ROOT_EVAL in data/coco_eval.py before running the evaluation.

    • PASCAL VOC: Use eval_voc.py.
    • MS-COCO: Use eval_coco.py.
    # Evaluation on PASCAL VOC
    python eval_voc.py --trained_model <trained weight path>
    
    # Evaluation on MS-COCO
    python eval_coco.py --trained_model <trained weight path>
  3. Prepare weights and configuration for training

    main

    Before training, you must download the backbone weights and configure your data paths:

    1. Create a weights directory and download the FC-reduced VGG-16 backbone weight into it.
    2. Update VOC_ROOT in data/voc0712.py or COCO_ROOT in data/coco.py to point to your dataset locations.
    3. Review data/config.py for training configurations.
    mkdir weights
    cd weights
    # Download vgg16_reducedfc.pth here
    cd ..
  4. Visualize model predictions with demo.py

    main

    Run the demo.py script to visualize the results of your trained model. You must provide the path to the trained weights via the --trained_model flag.

    python demo.py --trained_model <trained weight path>