SimCLR
SimCLR copied to clipboard
Similarity matrix shape does not match the shape of the mask
Hello,
I was currently testing the implementation when an error occured: The shape of the mask [512, 512] at index 0 does not match the shape of the indexed tensor [2, 2] at index 0. My batch size is 256.
The error occurs in this part of the code:
similarity_matrix = torch.matmul(features, features.T)
mask = torch.eye(labels.shape[0], dtype=torch.bool).to(device)
labels = labels[~mask].view(labels.shape[0], -1)
similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1)
I'm wondering if this something I'm doing wrong and how do I match the shape of tensors?
Thanks in advance!
same question~
The issue is because of the following line of code https://github.com/sthalles/SimCLR/blob/1848fc934ad844ae630e6c452300433fe99acfd9/simclr.py#L28 which assumes that the number of features you have is a multiple of the batch size but this is not always true.
For example, consider a dataset with only a 100 elements and a batch size of 256. In that case it will create a labels of size (256, 256) even though it should only be (200, 200) (assuming you are using n_views of 2). The way to resolve this is by updating the above line to:
labels = torch.cat([torch.arange(int(features.size(0)/2)) for i in range(self.args.n_views)], dim=0)