RuntimeError: The size of tensor a (3) must match the size of tensor b (0) at non-singleton dimension 0

Xhulio_Isufi: <p>Hi there, I want to implement semi-supervised learning. First, I want to do the slef-supervised part using SwAV from pl_bolts. I have my unlabelled dataset in a file in google drive and I am using google colab pro plus to run the code. Below you can find my code<br>
!pip install --upgrade pip<br>
!pip install pytorch-lightning==1.8.0<br>
!pip install git+https://github.com/PytorchLightning/lightning-bolts.git@master --upgrade<br>
from google.colab import drive<br>
drive.mount(‘/content/drive’)<br>
import os<br>
import torch<br>
import pytorch_lightning as pl<br>
from PIL import Image<br>
from torchvision import transforms<br>
from torch.utils.data import DataLoader<br>
from pl_bolts.models.self_supervised import SwAV<br>
from sklearn.metrics import accuracy_score<br>
from torchvision.transforms import Compose, RandomResizedCrop, RandomHorizontalFlip, Resize, ToTensor, Normalize<br>
import torch<br>
import torch.distributed as dist<br>
import torch.nn as nn<br>
import torchvision.transforms as transforms<br>
from torch.utils.data import DataLoader<br>
import pytorch_lightning as pl<br>
import os<br>
from PIL import Image<br>
import torch<br>
from torch.utils.data import Dataset<br>
class UnlabelledDataset(Dataset):<br>
def <strong>init</strong>(self, root, transform=None, num_views=2):<br>
self.root = root<br>
self.transform = transform<br>
self.num_views = num_views<br>
self.image_files = [f for f in os.listdir(root) if f.endswith((‘.png’, ‘.jpg’, ‘.jpeg’))]</p>
<pre><code>def __len__(self):
    return len(self.image_files)

def __getitem__(self, idx):
    image_path = os.path.join(self.root, self.image_files[idx])
    image = Image.open(image_path).convert('RGB')

    images = []
    for _ in range(self.num_views):
        transformed_image = image
        if self.transform:
            transformed_image = self.transform(image)
        images.append(transformed_image)

    return images, idx
</code></pre>
<p>def multi_view_collate_fn(batch):<br>
images, indices = zip(*batch)<br>
images = [torch.stack(view) for view in zip(*images)]<br>
images = torch.stack(images).transpose(0, 1).contiguous().view(-1, 3, 224, 224)<br>
indices = torch.tensor(indices)</p>
<pre><code>return images, indices
</code></pre>
<p>def train_swav(swav_checkpoint_path, unlabelled_root, num_epochs, batch_size=8):</p>
<pre><code>unlabelled_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.8, 0.8, 0.8, 0.2),
    transforms.RandomGrayscale(p=0.2),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])

'''unlabelled_dataset = UnlabelledDataset(root=unlabelled_root, transform=unlabelled_transforms)
unlabelled_loader = DataLoader(unlabelled_dataset, batch_size=batch_size, shuffle=True, num_workers=2, collate_fn=multi_view_collate_fn)


# Load the checkpoint and remove the process group from the state dict
checkpoint = torch.hub.load_state_dict_from_url(swav_checkpoint_path, map_location="cuda:0")
for key in list(checkpoint['state_dict'].keys()):
    if 'queue' in key:
        del checkpoint['state_dict'][key]

swav_model = SwAV(
    gpus=1,
    num_samples=len(unlabelled_dataset),
    batch_size=batch_size * 2,  # Account for the two views per image
    dataset='imagenet',
    nmb_crops=[2]  # Add this line to specify the number of views per image
)
swav_model.load_state_dict(checkpoint['state_dict'])
torch.autograd.set_detect_anomaly(True)
# Override the lr_scheduler_step method
def custom_lr_scheduler_step(self, epoch, val_loss, unused_arg):
    _, schedulers = self.configure_optimizers()
    for scheduler_dict in schedulers:
        scheduler_dict['scheduler'].step()

swav_model.lr_scheduler_step = custom_lr_scheduler_step.__get__(swav_model, SwAV)

trainer = pl.Trainer(gpus=1, max_epochs=num_epochs)
trainer.fit(swav_model, unlabelled_loader)

return swav_model.model
'''
unlabelled_dataset = UnlabelledDataset(root=unlabelled_root, transform=unlabelled_transforms)
unlabelled_dataset = torch.utils.data.Subset(unlabelled_dataset, range(0, 16))  # Use only the first 100 samples
unlabelled_loader = DataLoader(unlabelled_dataset, batch_size=batch_size, shuffle=True, num_workers=2, collate_fn=multi_view_collate_fn)

checkpoint = torch.hub.load_state_dict_from_url(swav_checkpoint_path, map_location="cuda:0")

swav_model = SwAV(
    gpus=1,
    dataset='imagenet',
    num_samples=len(unlabelled_dataset),
    batch_size=batch_size * 2,  # Account for the two views per image
    nmb_crops=[2, 224]  # two crops of size 224
)
swav_model.load_state_dict(checkpoint['state_dict'])
swav_model.size()
torch.autograd.set_detect_anomaly(True)

trainer = pl.Trainer(gpus=1, max_epochs=num_epochs)
trainer.fit(swav_model, unlabelled_loader)

return swav_model.model
</code></pre>
<p>pretrained_backbone = train_swav(‘<a href="https://pl-bolts-weights.s3.us-east-2.amazonaws.com/swav/swav_imagenet/swav_imagenet.pth.tar" rel="noopener nofollow ugc">https://pl-bolts-weights.s3.us-east-2.amazonaws.com/swav/swav_imagenet/swav_imagenet.pth.tar</a>’, ‘/content/drive/MyDrive/ThesisData/unlabeled/train05/’, num_epochs=1)</p>
<p>ANY HELP WILL BE MUCH APPRECIATED!</p>

Xhulio_Isufi: <p>Here is the full Traceback of the error</p>
<p>RuntimeError                              Traceback (most recent call last)<br>
 in &lt;cell line: 1&gt;()<br>
----&gt; 1 pretrained_backbone = train_swav(‘<a href="https://pl-bolts-weights.s3.us-east-2.amazonaws.com/swav/swav_imagenet/swav_imagenet.pth.tar" rel="noopener nofollow ugc">https://pl-bolts-weights.s3.us-east-2.amazonaws.com/swav/swav_imagenet/swav_imagenet.pth.tar</a>’, ‘/content/drive/MyDrive/ThesisData/unlabeled/train05/’, num_epochs=1)</p>
<p>35 frames<br>
 in train_swav(swav_checkpoint_path, unlabelled_root, num_epochs, batch_size)<br>
59<br>
60     trainer = pl.Trainer(gpus=1, max_epochs=num_epochs)<br>
—&gt; 61     trainer.fit(swav_model, unlabelled_loader)<br>
62<br>
63     return swav_model.model</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)<br>
606         model = self._maybe_unwrap_optimized(model)<br>
607         self.strategy._lightning_module = model<br>
 → 608         call._call_and_handle_interrupt(<br>
609             self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path<br>
610         )</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)<br>
36             return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)<br>
37         else:<br>
—&gt; 38             return trainer_fn(*args, **kwargs)<br>
39<br>
40     except _TunerExitException:</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py in _fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)<br>
648             model_connected=self.lightning_module is not None,<br>
649         )<br>
 → 650         self._run(model, ckpt_path=self.ckpt_path)<br>
651<br>
652         assert self.state.stopped</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py in _run(self, model, ckpt_path)<br>
1110         self._checkpoint_connector.resume_end()<br>
1111<br>
 → 1112         results = self._run_stage()<br>
1113<br>
1114         log.detail(f"{self.<strong>class</strong>.<strong>name</strong>}: trainer tearing down")</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py in _run_stage(self)<br>
1189         if self.predicting:<br>
1190             return self._run_predict()<br>
 → 1191         self._run_train()<br>
1192<br>
1193     def _pre_training_routine(self) → None:</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py in _run_train(self)<br>
1212<br>
1213         with torch.autograd.set_detect_anomaly(self._detect_anomaly):<br>
 → 1214             self.fit_loop.run()<br>
1215<br>
1216     def _run_evaluate(self) → _EVALUATE_OUTPUT:</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/loop.py in run(self, *args, **kwargs)<br>
197             try:<br>
198                 self.on_advance_start(*args, **kwargs)<br>
 → 199                 self.advance(*args, **kwargs)<br>
200                 self.on_advance_end()<br>
201                 self._restarting = False</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/fit_loop.py in advance(self)<br>
265         self._data_fetcher.setup(dataloader, batch_to_device=batch_to_device)<br>
266         with self.trainer.profiler.profile(“run_training_epoch”):<br>
 → 267             self._outputs = self.epoch_loop.run(self._data_fetcher)<br>
268<br>
269     def on_advance_end(self) → None:</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/loop.py in run(self, *args, **kwargs)<br>
197             try:<br>
198                 self.on_advance_start(*args, **kwargs)<br>
 → 199                 self.advance(*args, **kwargs)<br>
200                 self.on_advance_end()<br>
201                 self._restarting = False</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py in advance(self, data_fetcher)<br>
211<br>
212             with self.trainer.profiler.profile(“run_training_batch”):<br>
 → 213                 batch_output = self.batch_loop.run(kwargs)<br>
214<br>
215         self.batch_progress.increment_processed()</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/loop.py in run(self, *args, **kwargs)<br>
197             try:<br>
198                 self.on_advance_start(*args, **kwargs)<br>
 → 199                 self.advance(*args, **kwargs)<br>
200                 self.on_advance_end()<br>
201                 self._restarting = False</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/batch/training_batch_loop.py in advance(self, kwargs)<br>
86                 self.trainer.optimizers, self.trainer.optimizer_frequencies, kwargs.get(“batch_idx”, 0)<br>
87             )<br>
—&gt; 88             outputs = self.optimizer_loop.run(optimizers, kwargs)<br>
89         else:<br>
90             outputs = self.manual_loop.run(kwargs)</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/loop.py in run(self, *args, **kwargs)<br>
197             try:<br>
198                 self.on_advance_start(*args, **kwargs)<br>
 → 199                 self.advance(*args, **kwargs)<br>
200                 self.on_advance_end()<br>
201                 self._restarting = False</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py in advance(self, optimizers, kwargs)<br>
200         kwargs = self._build_kwargs(kwargs, self.optimizer_idx, self._hiddens)<br>
201<br>
 → 202         result = self._run_optimization(kwargs, self._optimizers[self.optim_progress.optimizer_position])<br>
203         if result.loss is not None:<br>
204             # automatic optimization assumes a loss needs to be returned for extras to be considered as the batch</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py in _run_optimization(self, kwargs, optimizer)<br>
247         else:<br>
248             # the <code>batch_idx</code> is optional with inter-batch parallelism<br>
 → 249             self._optimizer_step(optimizer, opt_idx, kwargs.get(“batch_idx”, 0), closure)<br>
250<br>
251         result = closure.consume_result()</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py in _optimizer_step(self, optimizer, opt_idx, batch_idx, train_step_and_backward_closure)<br>
368             )<br>
369             kwargs[“using_native_amp”] = isinstance(self.trainer.precision_plugin, MixedPrecisionPlugin)<br>
 → 370         self.trainer._call_lightning_module_hook(<br>
371             “optimizer_step”,<br>
372             self.trainer.current_epoch,</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py in _call_lightning_module_hook(self, hook_name, pl_module, *args, **kwargs)<br>
1354<br>
1355         with self.profiler.profile(f"[LightningModule]{pl_module.<strong>class</strong>.<strong>name</strong>}.{hook_name}"):<br>
 → 1356             output = fn(*args, **kwargs)<br>
1357<br>
1358         # restore current_fx when nested context</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/core/module.py in optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure, on_tpu, using_lbfgs)<br>
1752<br>
1753         “”"<br>
 → 1754         optimizer.step(closure=optimizer_closure)<br>
1755<br>
1756     def optimizer_zero_grad(self, epoch: int, batch_idx: int, optimizer: Optimizer, optimizer_idx: int) → None:</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/core/optimizer.py in step(self, closure, **kwargs)<br>
167<br>
168         assert self._strategy is not None<br>
 → 169         step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)<br>
170<br>
171         self._on_after_step()</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/strategy.py in optimizer_step(self, optimizer, opt_idx, closure, model, **kwargs)<br>
232         # TODO(fabric): remove assertion once strategy’s optimizer_step typing is fixed<br>
233         assert isinstance(model, pl.LightningModule)<br>
 → 234         return self.precision_plugin.optimizer_step(<br>
235             optimizer, model=model, optimizer_idx=opt_idx, closure=closure, **kwargs<br>
236         )</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/plugins/precision/precision_plugin.py in optimizer_step(self, optimizer, model, optimizer_idx, closure, **kwargs)<br>
117         “”“Hook to run the optimizer step.”“”<br>
118         closure = partial(self._wrap_closure, model, optimizer, optimizer_idx, closure)<br>
 → 119         return optimizer.step(closure=closure, **kwargs)<br>
120<br>
121     def _track_grad_norm(self, trainer: “pl.Trainer”) → None:</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/optim/lr_scheduler.py in wrapper(*args, **kwargs)<br>
67                 instance._step_count += 1<br>
68                 wrapped = func.<strong>get</strong>(instance, cls)<br>
—&gt; 69                 return wrapped(*args, **kwargs)<br>
70<br>
71             # Note that the returned function here is no longer a bound method,</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/optim/optimizer.py in wrapper(*args, **kwargs)<br>
278                                                f"but got {result}.")<br>
279<br>
 → 280                 out = func(*args, **kwargs)<br>
281                 self._optimizer_step_code()<br>
282</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/optim/optimizer.py in _use_grad(self, *args, **kwargs)<br>
31         try:<br>
32             torch.set_grad_enabled(self.defaults[‘differentiable’])<br>
—&gt; 33             ret = func(self, *args, **kwargs)<br>
34         finally:<br>
35             torch.set_grad_enabled(prev_grad)</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/optim/adam.py in step(self, closure)<br>
119         if closure is not None:<br>
120             with torch.enable_grad():<br>
 → 121                 loss = closure()<br>
122<br>
123         for group in self.param_groups:</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/plugins/precision/precision_plugin.py in _wrap_closure(self, model, optimizer, optimizer_idx, closure)<br>
103         consistent with the <code>PrecisionPlugin</code> subclasses that cannot pass <code>optimizer.step(closure)</code> directly.<br>
104         “”"<br>
 → 105         closure_result = closure()<br>
106         self._after_closure(model, optimizer, optimizer_idx)<br>
107         return closure_result</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py in <strong>call</strong>(self, *args, **kwargs)<br>
147<br>
148     def <strong>call</strong>(self, *args: Any, **kwargs: Any) → Optional[Tensor]:<br>
 → 149         self._result = self.closure(*args, **kwargs)<br>
150         return self._result.loss<br>
151</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py in closure(self, *args, **kwargs)<br>
133<br>
134     def closure(self, *args: Any, **kwargs: Any) → ClosureResult:<br>
 → 135         step_output = self._step_fn()<br>
136<br>
137         if step_output.closure_loss is None:</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/optimization/optimizer_loop.py in _training_step(self, kwargs)<br>
417         “”"<br>
418         # manually capture logged metrics<br>
 → 419         training_step_output = self.trainer._call_strategy_hook(“training_step”, *kwargs.values())<br>
420         self.trainer.strategy.post_training_step()<br>
421</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py in _call_strategy_hook(self, hook_name, *args, **kwargs)<br>
1492<br>
1493         with self.profiler.profile(f"[Strategy]{self.strategy.<strong>class</strong>.<strong>name</strong>}.{hook_name}"):<br>
 → 1494             output = fn(*args, **kwargs)<br>
1495<br>
1496         # restore current_fx when nested context</p>
<p>/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/strategy.py in training_step(self, *args, **kwargs)<br>
376         with self.precision_plugin.train_step_context():<br>
377             assert isinstance(self.model, TrainingStep)<br>
 → 378             return self.model.training_step(*args, **kwargs)<br>
379<br>
380     def post_training_step(self) → None:</p>
<p>/usr/local/lib/python3.10/dist-packages/pl_bolts/models/self_supervised/swav/swav_module.py in training_step(self, batch, batch_idx)<br>
230<br>
231     def training_step(self, batch, batch_idx):<br>
 → 232         loss = self.shared_step(batch)<br>
233<br>
234         self.log(“train_loss”, loss, on_step=True, on_epoch=False)</p>
<p>/usr/local/lib/python3.10/dist-packages/pl_bolts/models/self_supervised/swav/swav_module.py in shared_step(self, batch)<br>
217<br>
218         # SWAV loss computation<br>
 → 219         loss, queue, use_queue = self.criterion(<br>
220             output=output,<br>
221             embedding=embedding,</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)<br>
1499                 or _global_backward_pre_hooks or _global_backward_hooks<br>
1500                 or _global_forward_hooks or _global_forward_pre_hooks):<br>
 → 1501             return forward_call(*args, **kwargs)<br>
1502         # Do not call functions when jit is used<br>
1503         full_backward_hooks, non_full_backward_hooks = [], []</p>
<p>/usr/local/lib/python3.10/dist-packages/pl_bolts/models/self_supervised/swav/loss.py in forward(self, output, embedding, prototype_weights, batch_size, queue, use_queue)<br>
74             for v in np.delete(np.arange(np.sum(self.num_crops)), crop_id):<br>
75                 p = self.softmax(output[batch_size * v : batch_size * (v + 1)] / self.temperature)<br>
—&gt; 76                 subloss -= torch.mean(torch.sum(q * torch.log(p), dim=1))<br>
77             loss += subloss / (np.sum(self.num_crops) - 1)<br>
78         loss /= len(self.crops_for_assign)  # type: ignore</p>
<p>RuntimeError: The size of tensor a (3) must match the size of tensor b (0) at non-singleton dimension 0</p>

ptrblck: <p>Based on the error message and stacktrace the loss calculation fails in:</p>
<pre><code class="lang-python">subloss -= torch.mean(torch.sum(q * torch.log(p), dim=1))
</code></pre>
<p>with</p>
<pre><code class="lang-python">RuntimeError: The size of tensor a (3) must match the size of tensor b (0) at non-singleton dimension 0
</code></pre>
<p>Check the shapes of both tensors (<code>q</code> and <code>p</code>) and make sure you can properly multiply them as the shape mismatch causes the issue currently.</p>