HobbitLong / SupContrast

PyTorch implementation of "Supervised Contrastive Learning" (and SimCLR incidentally)
BSD 2-Clause "Simplified" License
2.98k stars 525 forks source link

Provided pre-trained load issue #115

Open Rasoul77 opened 1 year ago

Rasoul77 commented 1 year ago

I tried to use the provided pre-trained PTH file presented in the Update section using the following code snippet, but it returns errors of missing keys,

import torch
from networks.resnet_big import SupConResNet

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"device to use: {device}")

model = SupConResNet()
model.to(device=device)
model.load_state_dict(torch.load("./supcon.pth", map_location=device))
model.eval()

which results,

Error(s) in loading state_dict for SupConResNet:
    Missing key(s) in state_dict: "encoder.conv1.weight", "encoder.bn1.weight", ...

Could you clarify for which model the pre-trained "supcon.pth" is provided?

FraCorti commented 1 year ago

Hi @Rasoul77, by chance did you find a solution?

FraCorti commented 1 year ago

@Rasoul77 I solved the serialization by creating the ResNet50 model from the old implementation (https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py), changing the first layer kernel size to 7 and deleting the fully connected part and taking only the projection head. The class that represents the network is SupConResNet.

Khey17 commented 1 year ago

@Rasoul77 Did you find a solution? I got the same error when i ran main_linear.py. image