How to get image embedding with vicreg?

[MikeMACintosh]

Hi, thanks for your work!
How can i get image features from vicreg models?
I doing:

import torch 
vicreg_model = torch.hub.load('facebookresearch/vicreg:main', 'resnet50')
data_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

data = torch.unsqueeze(data_transforms(Image.open(image).convert('RGB')), dim=0)
with torch.no_grad():  
    embedding = vicreg_model(data.float()).cpu()

And question for further works: how i can implement self-supervised metric learning with vicreg for custom dataset? I need a good embedding for image from my own dataset.

[Adrien987k]

Hi,

Do you have an issue for running the code above ?

For further work you can just take the code, replace the dataloader by your custom dataloader, and pretrain on your dataset. You can also fine-tune the pretrained model on your custom dataset.