DOCSEARCHbyAlgolia
Submit Docs

MixVPR: All-MLP Feature Aggregation for Visual Place Recognition

repository·main·Indexed 43 days ago

https://github.com/amaralibey/mixvpr

An all-MLP feature aggregation method for large-scale Visual Place Recognition (VPR) designed for high-performance and low-latency. Built with PyTorch Lightning, it leverages pre-trained backbones like ResNet50 and supports various aggregators including MixVPR, ConvAP, and NetVLAD. The VPRModel class handles feature extraction, online mining with MultiSimilarity or TripletMargin miners, and linear learning rate warmup.

Tokens
1.6K
Snippets
3
Records
6
Agent score
6%

What's inside MixVPR

  1. Load pretrained MixVPR models

    main

    To use MixVPR with pretrained weights, instantiate a VPRModel from the main module and load a state dictionary using torch.load.

    Important Requirement: All input images must be resized to 320x320 pixels before being passed to the model.

    Available ResNet50 backbones include output dimensions of 4096, 512, and 128. Ensure your agg_config matches the specific architecture of the weights you are loading.

    from main import VPRModel
import torch

# Note that images must be resized to 320x320
model = VPRModel(backbone_arch='resnet50', 
                 layers_to_crop=[4],
                 agg_arch='MixVPR',
                 agg_config={'in_channels' : 1024,
                             'in_h' : 20,
                             'in_w' : 20,
                             'out_channels' : 1024,
                             'mix_depth' : 4,
                             'mlp_ratio' : 1,
                             'out_rows' : 4},
                )

state_dict = torch.load('./LOGS/resnet50_MixVPR_4096_channels(1024)_rows(4).ckpt')
model.load_state_dict(state_dict)
model.eval()

  2. Understand the learning rate warmup mechanism

    main

    The model implements a custom optimizer_step to handle a linear learning rate warmup phase. If the current global step is less than warmpup_steps, the learning rate is scaled linearly from 0 up to the target lr.

    Formula: lr_scale = min(1., (global_step + 1) / warmpup_steps) current_lr = lr_scale * target_lr

  3. Initialize the VPRModel class

    main

    The VPRModel class is the primary entry point for Visual Place Recognition tasks. It is built using PyTorch Lightning and consists of a backbone (encoder) and an aggregator.

    Key configuration parameters include:

    • Backbone Configuration:
      • backbone_arch: The architecture name (e.g., 'resnet50', 'resnet101', 'efficientnet_b0').
      • pretrained: Boolean indicating if pretrained weights should be used.
      • layers_to_freeze: Number of layers to freeze.
      • layers_to_crop: List of indices for cropping layers (e.g., [4] to crop the last ResNet layer).
    • Aggregator Configuration:
      • agg_arch: The aggregator architecture (e.g., 'MixVPR', 'ConvAP', 'CosPlace', 'NetVLAD', 'GeM').
      • agg_config: A dictionary of parameters specific to the chosen aggregator.
    • Training Hyperparameters:
      • lr: Learning rate.
      • optimizer: 'sgd', 'adamw', or 'adam'.
      • warmpup_steps: Number of steps for the learning rate warmup phase.
      • milestones: List of epochs for the MultiStepLR scheduler.
      • lr_mult: Multiplier for the learning rate decay.
    • Loss and Mining:
      • loss_name: The loss function to use (e.g., 'MultiSimilarityLoss', 'ContrastiveLoss', 'TripletMarginLoss').
      • miner_name: The online miner to use (e.g., 'MultiSimilarityMiner', 'TripletMarginMiner').
      • miner_margin: Margin for the miner.
    from main import VPRModel

model = VPRModel(
    backbone_arch='resnet50',
    pretrained=True,
    layers_to_freeze=2,
    layers_to_crop=[4],
    agg_arch='MixVPR',
    agg_config={
        'in_channels': 1024,
        'in_h': 20,
        'in_w': 20,
        'out_channels': 1024,
        'mix_depth': 4,
        'mlp_ratio': 1,
        'out_rows': 4
    },
    lr=0.05,
    optimizer='sgd',
    loss_name='MultiSimilarityLoss',
    miner_name='MultiSimilarityMiner'
)

  4. Configure the loss function and online mining

    main

    The loss_function(descriptors, labels) method handles the calculation of the training loss. It supports two modes:

    1. With Online Mining: If a miner_name is provided, the model uses the specified miner to select hard pairs/triplets from the batch before calculating the loss. The model also tracks b_acc (batch accuracy), which represents the percentage of samples that contributed to the loss (i.e., were not trivial).
    2. Without Online Mining: The loss is calculated directly from the descriptors and labels.

    If a loss function returns a tuple (e.g., (loss, batch_acc)), the model will use that tuple to log the batch accuracy.

  5. Use VPRModel for forward passes

    main

    The VPRModel.forward(x) method performs the full feature extraction pipeline: it passes the input tensor x through the backbone to extract features, and then through the aggregator to produce descriptors.

    Note: Images must be resized to 320x320 before being passed to the model.

    # Assuming model is an instance of VPRModel
# x should be a tensor of shape (Batch, Channels, Height, Width)
# Height and Width should be 320
descriptors = model(x)

  6. Configure optimizers and learning rate scheduling

    main

    The VPRModel implements configure_optimizers to support several optimization strategies and a learning rate scheduler.

    Supported Optimizers:

    • sgd: Uses torch.optim.SGD with lr, weight_decay, and momentum.
    • adamw: Uses torch.optim.AdamW with lr and weight_decay.
    • adam: Uses torch.optim.AdamW with lr and weight_decay.

    Scheduler: A MultiStepLR scheduler is automatically applied using the milestones and lr_mult parameters provided during initialization.