Note: We will get combined VOC2007 and VOC2012 dataset when we will use name 'voc' in the wrapper functions.

• One can get different data splits (train, and test splits) using the wrapper function get_dataloaders,

    data_splits = get_dataloaders(
        dataset_name='voc', batch_size=64
    )
    train_split = data_splits['train']
    test_split = data_splits['test']

• To get the desired model architecture, we have a wrapper function get_model, which requires exact model and dataset name.

    model = get_model(
        model_name="yolo-v5s",
        dataset_name="voc",
        pretrained=True, # or False, if pretrained weights are not required
    )

• To get the evaluation function, we have a wrapper function get_eval_by_name, which requires exact model and dataset name.

    test_loader = data_splits["test"]
    eval_function = get_eval_function(dataset_name="voc", model_name="yolo-v5s")
    APs = eval_function(model, test_loader)

List of models and corresponding datasets used to train it can be found here.

One could create a model with a custom number of classes, while loading the pretrained weights from one of the pretrained models available in the zoo for the modules where weight reusage is possible. An example below creates a yolo5_6m model with 8 output classes, with weights loaded from a COCO checkpoint (except for the layers whise shape depends on the number of classes):

    from deeplite_torch_zoo import get_model

    model = get_model(
        model_name="yolo5_6m",
        dataset_name="coco",
        num_classes=8,
        pretrained=True, # or False, if pretrained weights are not required
    )

It needs to be ensured that the data format should follow either of the formats present in the available datasets with proper names: voc, coco, coco_person. One can train any yolo model with resume functionalities using train_yolo.py. Command to run this script,

    $ python train_yolo.py
            --dataset ZOO_DATASET_NAME
            --img-dir DATA_IMG_ROOT_PATH
            --net ZOO_MODEL_NAME
            --hp_config scratch                 # Options: 'scratch', 'finetune'
            --test_img_res 224
            --train_img_res 224
            --num-classes NUM_CLASSES
            --weight_path WEIGHT_PATH
            --batch-size BATCH_SIZE
            --epochs NUM_EPOCHS
            --device 0

Similarly for training any SSD model, one can use train_ssd.py with proper arguments. Also, to access the pretrained weights, one can add --pretrained in the arguments along with the --pretraining_source_dataset which will again be selected from the available options.

List of all models with their names and corresponding datasets used to train it.