CVNets Computer Vision Toolkit

repository·main·Indexed 24 days ago

https://github.com/apple/ml-cvnets

A computer vision toolkit for training mobile-optimized and non-mobile models. CVNets supports tasks including image classification (CNNs, ViT, MobileViT, SwinTransformer), object detection (SSD, Mask R-CNN), semantic segmentation (DeepLabv3, PSPNet), multimodal classification (ByteFormer), and foundation models (CLIP). It also features automatic data augmentation via RangeAugment, as well as soft and hard distillation.

Tokens
35.9K
Snippets
70
Records
95
Agent score
83%

What's inside CVNets

  1. Overview of CVNets

    main

    CVNets is an open-source library designed for training deep neural networks for visual recognition tasks. It supports classification, detection, and segmentation. The library provides tools for image and video understanding, including:

    • Data loading and transformations
    • Novel data sampling methods
    • Implementations of several state-of-the-art networks
  2. Use RangeAugment for various computer vision tasks

    main

    RangeAugment (referred to as Neural Augmentor or NA in the cvnets/neural_augmentor codebase) is an automatic augmentation method that learns model- and task-specific magnitude ranges for augmentation operations.

    CVNets provides training and evaluation code, pretrained models, and configuration files for the following tasks:

    • Image Classification: On the ImageNet dataset.
    • Semantic Segmentation: On the ADE20k and PASCAL VOC datasets.
    • Object Detection: On the MS-COCO dataset.
    • Contrastive Learning: Using Image-Text pairs (CLIP style).
    • Distillation: On the ImageNet dataset.
  3. Understand the CVNets directory structure

    main

    The CVNets codebase is organized into functional areas. Key directories for users include:

    • cvnets/models/: Task-specific model definitions (classification, detection, segmentation, etc.).
    • cvnets/modules/: Reusable high-level building blocks (e.g., InvertedResidual, TransformerEncoder).
    • data/datasets/: Task-specific dataset classes (e.g., ImagenetDataset).
    • data/transforms/: Image and video transformations. PIL-based transforms are recommended for better performance.
    • data/sampler/: Custom data samplers.
    • engine/: Core training and evaluation logic (training_engine.py and evaluation_engine.py).
  4. Supported models and tasks in CVNets

    main

    CVNets supports a wide range of computer vision models and tasks. You can find a full list in the Model Zoo or the examples folder.

    Key supported categories include:

    • ImageNet Classification:
      • CNNs: MobileNet (v1, v2, v3), EfficientNet, ResNet, RegNet.
      • Transformers: Vision Transformer (ViT), MobileViT (v1, v2), SwinTransformer.
    • Multimodal Classification: ByteFormer.
    • Object Detection: SSD, Mask R-CNN.
    • Semantic Segmentation: DeepLabv3, PSPNet.
    • Foundation Models: CLIP.
    • Automatic Data Augmentation: RangeAugment, AutoAugment, RandAugment.
    • Distillation: Soft distillation and Hard distillation.
  5. How the training and evaluation workflow works

    main

    The training and evaluation process follows a two-step pattern:

    1. Initialization: Entry scripts (main_train.py for training or main_eval.py for evaluation) are executed. These scripts are responsible for building and initializing the model, dataset, optimizer, and setting up distributed training if required.
    2. Execution: The initialized objects are passed to the Training/Evaluation Engine (engine/training_engine.py and engine/evaluation_engine.py), which contains the actual logic for running the training or evaluation loops.
  6. Use variably-sized video samplers

    main

    CVNet includes a specialized sampler for video data (video_variable_seq_sampler). This allows researchers to control video-specific input variables to learn space- and time-invariant representations. Key controllable variables include:

    • Number of frames
    • Number of clips per video
    • Video spatial resolution
  7. How models and tasks are organized in CVNets

    main

    Models are categorized by task under cvnets/models/<task>. Each task utilizes a specific parent class to ensure consistency. For example, classification models derive from cvnets.models.classification.base_cls.BaseEncoder.

    Models are designed to be reusable across different tasks. A model defined for classification (like ResNet) can be used as an encoder for a detection model (like ssd) to prevent code duplication.

  8. Understand CVNet data sampling strategies

    main

    CVNet provides three primary sampling strategies for training computer vision networks, which allow you to balance scale robustness against GPU memory utilization:

    1. Single-scale with fixed batch size (SSc-FBS): The standard approach where every batch uses a pre-defined spatial resolution $(H, W)$ and a fixed batch size $b$ per GPU.
    2. Multi-scale with fixed batch size (MSc-FBS): Samples from a sorted set of $n$ spatial resolutions $\mathcal{S} = {(H_1, W_1), \dots, (H_n, W_n)}$. At each iteration, a resolution is randomly selected from the set, but the batch size $b$ remains constant. This improves scale robustness but can lead to high peak GPU memory usage.
    3. Multi-scale with variable batch size (MSc-VBS): An extension of MSc-FBS designed to prevent Out-of-Memory (OOM) errors. While it still samples from a set of resolutions $\mathcal{S}$, it adjusts the batch size $b_t$ for a chosen resolution $(H_t, W_t)$ based on the maximum resolution $(H_n, W_n)$ and the target batch size $b$ for that maximum resolution. The formula used is: $b_t = \frac{H_n W_n b}{H_t W_t}$.
  9. Configure ByteFormer for accurate FLOPs and model size estimates

    main

    ByteFormer configurations often use larger embedding sizes than strictly necessary to allow for experimentation with different kernel sizes or input types. To get accurate performance estimates (FLOPs/Model Size) for your specific input domain, you must adjust the following parameters:

    1. --model.classification.byteformer.max-num-tokens: Set this to the average token length (after the Conv1D downsampling) for your specific input type (e.g., TIFF, JPEG).
      • Warning: Do not set this to the average length during training if you have variable-length inputs (like JPEG), as an input exceeding this value will cause an error.
    2. --model.classification.byteformer.dummy-input-token-length: Set this to your expected input length before the Conv1D downsampling for your particular input domain.
  10. Quickstart: Segment an image using DeepLabv3 with MobileViT

    main

    This example demonstrates segmenting an image from a URL using a DeepLabv3 model with a MobileViT backbone.

    export IMG_PATH="http://farm7.staticflickr.com/6206/6118204766_b1c9a39153_z.jpg"
    export CFG_FILE="https://docs-assets.developer.apple.com/ml-research/models/cvnets-v2/segmentation/pascalvoc/deeplabv3-mobilevitv1.yaml"
    export MODEL_WEIGHTS="https://docs-assets.developer.apple.com/ml-research/models/cvnets-v2/segmentation/pascalvoc/deeplabv3-mobilevitv1.pt"
    cvnets-eval-seg --common.config-file $CFG_FILE --common.results-loc deeplabv3_results --model.segmentation.pretrained $MODEL_WEIGHTS --model.segmentation.n-classes 21 \
     --evaluation.segmentation.resize-input-images --evaluation.segmentation.mode single_image --evaluation.segmentation.path "${IMG_PATH}" --evaluation.segmentation.save-masks \
     --evaluation.segmentation.apply-color-map --evaluation.segmentation.save-overlay-rgb-pred
  11. Register a new dataset type

    main

    To add a new dataset to CVNets, you must register your dataset class with data.dataset.DATASET_REGISTRY. Use the @DATASET_REGISTRY.register decorator to specify a name and a type (task). This registration allows the dataset to be instantiated via configuration files.

    Datasets should ideally inherit from BaseImageDataset or BaseVideoDataset (which both inherit from BaseDataset), though this is currently a soft requirement.

    from data.datasets import DATASET_REGISTRY
    from data.datasets.dataset_base import BaseImageDataset
    
    @DATASET_REGISTRY.register(name="ade20k", type="segmentation")
    class ADE20KDataset(BaseImageDataset):
        # PyTorch Dataset implementation
        pass