DOCSEARCHbyAlgolia
Submit Docs

WTConv

repository·main·Indexed 20 days ago

https://github.com/bgu-cs-vil/wtconv

Implementation of Wavelet Convolutions to achieve large receptive fields in CNNs. Includes standard PyTorch layers and Fast WTConv, which provides optimized backends for CUDA, Metal (MPS), and Triton. Features WTConvNeXt models registered with the timm library and integration guides for training on ImageNet-1k and semantic segmentation via MMSegmentation.

Tokens
1.9K
Snippets
7
Records
12
Agent score
70%

What's inside wtconv

  1. Overview of Fast WTConv Backends

    main

    Fast WTConv provides optimized backends for different hardware environments:

    • CUDA: Optimized kernels for NVIDIA GPUs. Supports fp32, fp16, and bf16.
    • Metal (MPS): Optimized shaders for Apple Silicon (M1/M2/M3). Supports fp32 and fp16.
    • Triton: Pure Triton implementation for portability and high performance without CUDA dependencies. Supports fp32, fp16, and bf16.

  2. Use Fast WTConv for high performance

    main
    For high-performance requirements, use fast_wtconv. It provides optimized implementations for CUDA, Metal (MPS), and Triton backends. Refer to the fast_wtconv/README.md file for specific installation and usage details.

  3. Set up Semantic Segmentation with WTConvNeXt

    main

    To use WTConvNeXt for semantic segmentation, you must integrate it into the MMSegmentation framework.

    Prerequisites

    • MMSegmentation: Follow the official installation guide. Tested with mmsegmentation==2.2.0, python==3.9, pytorch==1.9.1, and cuda 11.1.
    • MMPretrain: Install mmpretrain (formerly MMClassification) via pip:

    Integration Steps

    1. File Placement: Place the WTConvNeXt configuration files and model files into the corresponding folders within your MMSegmentation directory.
    2. Register Backbone: You must manually register the backbone in MMSegmentation:
      • Open mmseg/models/backbones/__init__.py.
      • Add from .wtconvnext import WTConvNeXt.
      • Add WTConvNeXt to the list of model names in that file.
    3. Model Key Translation: Because the model keys in the provided save files differ from MMSegmentation's expected keys, run the provided translation script to prepare the pretrained weights:
      • Use translate_model_to_mmseg.py on your pretrained save file.
    pip install mmpretrain>=1.0.0

  4. Evaluate WTConvNeXt on ImageNet-1k

    main

    To validate a trained WTConvNeXt model, use the validate.py script (adapted from timm).

    python validate.py --model wtconvnext_tiny \
                   --data-dir IMAGENET_PATH \
                   --checkpoint WTConvNeXt_tiny_5_300e_ema.pth

  5. Train WTConvNeXt on ImageNet-1k

    main

    Training scripts are adapted from the timm library. To maintain an effective batch size of 4096 when using multiple GPUs, adjust the batch-size and grad-accum-steps such that (gpus * batch-size * grad-accum-steps) = 4096.

    # Example: Multi-GPU training with 4 GPUs
torchrun --nproc-per-node=4  \/n         python train.py --model wtconvnext_tiny --drop-path 0.1 \/n                --data-dir IMAGENET_PATH \/n                --epochs 300 --warmup-epochs 20 \/n                --batch-size 64 --grad-accum-steps 16 --sched-on-updates \/n                --lr 4e-3 --weight-decay 5e-2 \/n                --opt adamw --layer-decay 1.0 \/n                --aa rand-m9-mstd0.5-inc1 \/n                --reprob 0.25 --mixup 0.8 --cutmix 1.0 \/n                --model-ema --model-ema-decay 0.9999 \/n                --output checkpoints/wtconvnext_tiny_300/

  6. Configure Batch Size for WTConvNeXt Training

    main

    When training WTConvNeXt, ensure you use an effective batch size of 16.

    The provided configuration files are tuned for a setup using 8 GPUs with a batch size of 2 per GPU ($8 \times 2 = 16$). If your hardware configuration differs, you must manually adjust the batch_size in the config file to maintain this effective batch size.

  7. Install Fast WTConv

    main

    To use fast_wtconv, ensure you have the following dependencies installed:

    • PyTorch
    • Triton (required if you intend to use the Triton backend)

    Note on CUDA Backend: If you use the CUDA backend, you must have nvcc (NVIDIA CUDA Compiler) installed and available in your system PATH because all implementations use JIT (Just-In-Time) compilation.

  9. Use WTConvNeXt via timm registry

    main

    WTConvNeXt models are registered with the timm library, allowing you to instantiate them using create_model.

    import wtconvnext
from timm.models import create_model

model = create_model(
    "wtconvnext_tiny",
    pretrained=False,
    num_classes=1000
)

  10. Use WTConv2d in a CNN

    main

    You can integrate wavelet convolution layers directly into your custom CNN architectures by importing WTConv2d from the wtconv package.

    from wtconv import WTConv2d

# Example: 32 input channels, 32 output channels, 5x5 kernel, 3 wavelet levels
conv_dw = WTConv2d(32, 32, kernel_size=5, wt_levels=3)

  11. Use WTConv2d with Auto-Backend Detection

    main

    The standard WTConv2d class from fast_wtconv.wtconv automatically detects your hardware (CUDA for NVIDIA GPUs or MPS for Apple Silicon) and selects the appropriate optimized kernel. This makes it a drop-in replacement for original WTConv layers.

    Initialization Parameters:

    • in_channels: Number of input channels.
    • out_channels: Number of output channels.
    • kernel_size: Size of the convolution kernel.
    • stride: Stride of the convolution.
    • wt_levels: Number of wavelet levels.
    import torch
from fast_wtconv.wtconv import WTConv2d

# Initialize layer
# Parameters: in_channels, out_channels, kernel_size, stride, wt_levels
layer = WTConv2d(64, 64, kernel_size=5, wt_levels=2)

# Move to device (CUDA or MPS)
device = 'cuda' if torch.cuda.is_available() else 'mps'
layer = layer.to(device)

# Forward pass
x = torch.randn(1, 64, 224, 224).to(device)
output = layer(x)

  12. Use the Triton Backend for WTConv

    main

    If you require a pure Triton implementation (for example, for AMD GPUs or specific performance profiles), you can use the WTConv2d class from the fast_wtconv.wtconv_triton module. This implementation is portable and does not depend on CUDA.

    import torch
from fast_wtconv.wtconv_triton import WTConv2d as WTConv2dTriton

# Initialize Triton layer
layer = WTConv2dTriton(64, 64, kernel_size=5, wt_levels=2).cuda()

# Forward pass
x = torch.randn(1, 64, 224, 224).cuda()
output = layer(x)