All In 1

yolo.lazy is a packaged file that includes training, validation, and inference tasks. For detailed function documentation, thercheck out the IPython notebooks to learn how to import and use these function the following section will break down operation inside of lazy, also supporting directly import/call the function.

[TOC], setup, build, dataset, train, validation, inference To train the model, you can run:

Train Model

  • batch size check / cuda

  • training time / check

  • build model / check

  • dataset / check

python yolo/lazy.py task=train

You can customize the training process by overriding the following common arguments:

  • name: str The experiment name.

  • model: str Model backbone, options include [model_zoo] v9-c, v7, v9-e, etc.

  • cpu_num: int Number of CPU workers (num_workers).

  • out_path: Path The output path for saving models and logs.

  • weight: Path | bool | None The path to pre-trained weights, False for training from scratch, None for default weights.

  • use_wandb: bool Whether to use Weights and Biases for experiment tracking.

  • use_TensorBoard: bool Whether to use TensorBoard for logging.

  • image_size: int | [int, int] The input image size.

  • +quiet: bool Optional, disable all output.

  • task.epoch: int Total number of training epochs.

  • task.data.batch_size: int The size of each batch (auto-batch sizing [WIP]).

Examples

To train a model with a specific batch size and image size, you can run:

python yolo/lazy.py task=train task.data.batch_size=12 image_size=1280

Multi-GPU Training with DDP

For multi-GPU training, we use Distributed Data Parallel (DDP) for efficient and scalable training. DDP enable training model with mutliple GPU, even the GPUs aren’t on the same machine. For more details, you can refer to the DDP tutorial.

To train on multiple GPUs, replace the python command with torchrun --nproc_per_node=[GPU_NUM]. The nproc_per_node argument specifies the number of GPUs to use.

torchrun --nproc_per_node=2 yolo/lazy.py task=train device=[0,1]

Training on a Custom Dataset

To use the auto-download module, we suggest users construct the dataset config in the following format. If the config files include auto_download, the model will automatically download the dataset when creating the dataloader.

Here is an example dataset config file:

path: data/dev
train: train
validation: val

class_num: 80
class_list: ['Person', 'Bicycle', 'Car', 'Motorcycle', 'Airplane', 'Bus', 'Train', 'Truck', 'Boat', 'Traffic light', 'Fire hydrant', 'Stop sign', 'Parking meter', 'Bench', 'Bird', 'Cat', 'Dog', 'Horse', 'Sheep', 'Cow', 'Elephant', 'Bear', 'Zebra', 'Giraffe', 'Backpack', 'Umbrella', 'Handbag', 'Tie', 'Suitcase', 'Frisbee', 'Skis', 'Snowboard', 'Sports ball', 'Kite', 'Baseball bat', 'Baseball glove', 'Skateboard', 'Surfboard', 'Tennis racket', 'Bottle', 'Wine glass', 'Cup', 'Fork', 'Knife', 'Spoon', 'Bowl', 'Banana', 'Apple', 'Sandwich', 'Orange', 'Broccoli', 'Carrot', 'Hot dog', 'Pizza', 'Donut', 'Cake', 'Chair', 'Couch', 'Potted plant', 'Bed', 'Dining table', 'Toilet', 'Tv', 'Laptop', 'Mouse', 'Remote', 'Keyboard', 'Cell phone', 'Microwave', 'Oven', 'Toaster', 'Sink', 'Refrigerator', 'Book', 'Clock', 'Vase', 'Scissors', 'Teddy bear', 'Hair drier', 'Toothbrush']

auto_download:

Both of the following formats are acceptable:

  • path: str The path to the dataset.

  • train, validation: str The training and validation directory names under /images. If using txt as ground truth, these should also be the names under /labels/.

  • class_num: int The number of dataset classes.

  • class_list: List[str] Optional, the list of class names, used only for visualizing the bounding box classes.

  • auto_download: dict Optional, whether to auto-download the dataset.

The dataset should include labels or annotations, preferably in JSON format for compatibility with pycocotools during inference:

DataSetName/
├── annotations
│   ├── train_json_name.json
│   └── val_json_name.json
├── labels/
│   ├── train/
│   │   ├── AnyLabelName.txt
│   │   └── ...
│   └── validation/
│       └── ...
└── images/
    ├── train/
    │   ├── AnyImageNameN.{png,jpg,jpeg}
    │   └── ...
    └── validation/
        └── ...

Validation Model

During training, this block will be auto-executed. You may also run this task manually to generate a JSON file representing the predictions for a given validation dataset. If the validation set includes JSON annotations, it will run pycocotools for evaluation.

We recommend setting task.data.shuffle to False and turning off task.data.data_augment.

You can customize the validation process by overriding the following arguments:

  • task.nms.min_confidence: str The minimum confidence of model prediction.

  • task.nms.min_iou: str The minimum IoU threshold for NMS (Non-Maximum Suppression).

Examples

python yolo/lazy.py task=validation task.nms.min_iou=0.9

Model Inference

Note

The dataset parameter shouldn’t be overridden because the model requires the class_num of the dataset. If the classes have names, please provide the class_list.

You can customize the inference process by overriding the following arguments:

  • task.fast_inference: str Optional. Values can be onnx, trt, deploy, or None. deploy will detach the model auxiliary head.

  • task.data.source: str | Path | int This argument will be auto-resolved and could be a webcam ID, image folder path, video/image path.

  • task.nms.min_confidence: str The minimum confidence of model prediction.

  • task.nms.min_iou: str The minimum IoU threshold for NMS (Non-Maximum Suppression).

Examples

python yolo/lazy.py model=v9-m task.nms.min_confidence=0.1 task.data.source=0 task.fast_inference=onnx