DOCSEARCHbyAlgolia
Submit Docs

VICReg Documentation

repository·main·Indexed 20 days ago

https://github.com/facebookresearch/vicreg

PyTorch implementation of Variance-Invariance-Covariance Regularization for self-supervised learning. Includes scripts for single-node and multi-node (SLURM) pretraining, linear and semi-supervised evaluation on ImageNet, and pretrained models such as resnet50, resnet50x2, and resnet200x2 available via PyTorch Hub.

Tokens
982
Snippets
5
Records
6
Agent score
19%

What's inside VICReg

  1. Pretrain VICReg on multiple nodes using SLURM

    main

    For multi-node training, use the run_with_submitit.py script. This requires the submitit package (pip install submitit).

    Example for 4 nodes with 8 GPUs each for 1000 epochs:

    python run_with_submitit.py --nodes 4 --ngpus 8 --data-dir /path/to/imagenet --exp-dir /path/to/experiment/ --arch resnet50 --epochs 1000 --batch-size 2048 --base-lr 0.2

  2. Load pretrained VICReg models via PyTorch Hub

    main

    You can easily load pretrained VICReg backbones directly into your PyTorch workflow using torch.hub.load. The available models include resnet50, resnet50x2, and resnet200x2.

    import torch
resnet50 = torch.hub.load('facebookresearch/vicreg:main', 'resnet50')
resnet50x2 = torch.hub.load('facebookresearch/vicreg:main', 'resnet50x2')
resnet200x2 = torch.hub.load('facebookresearch/vicreg:main', 'resnet200x2')

  3. Perform semi-supervised evaluation

    main

    To evaluate a pretrained model via semi-supervised fine-tuning on a subset of ImageNet labels, use evaluate.py with the --weights finetune flag. You can specify the percentage of labels used via --train-perc.

    Example for 1% label fine-tuning:

    python evaluate.py --data-dir /path/to/imagenet/ --pretrained /path/to/checkpoint/resnet50.pth --exp-dir /path/to/experiment/ --weights finetune --train-perc 1 --epochs 20 --lr-backbone 0.03 --lr-classifier 0.08 --weight-decay 0

  4. Perform linear evaluation on pretrained models

    main

    To evaluate a pretrained backbone using linear classification on ImageNet, use the evaluate.py script. Use the --pretrained flag to point to your checkpoint and --lr-head to set the learning rate for the linear head.

    python evaluate.py --data-dir /path/to/imagenet/ --pretrained /path/to/checkpoint/resnet50.pth --exp-dir /path/to/experiment/ --lr-head 0.02

  5. Pretrain VICReg on a single node

    main

    To perform single-node local training using PyTorch distributed launch, use the main_vicreg.py script. You must provide the ImageNet data directory, an experiment directory, the architecture, and training hyperparameters.

    Example for ResNet-50 with 8 GPUs for 100 epochs:

    python -m torch.distributed.launch --nproc_per_node=8 main_vicreg.py --data-dir /path/to/imagenet/ --exp-dir /path/to/experiment/ --arch resnet50 --epochs 100 --batch-size 512 --base-lr 0.3