Expected all tensors to be on the same devices

[DickyAdi]

This is my first time in object detection task and pytorch lightning also torchmetrics.

When computing the IoU on the validation step, sometimes in the first few epochs the code run fine without any error however after a few epoch there's an error saying "Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument tensors in method wrapper_cat)." How does that happen after the first few epochs run perfectly fine?

Here i provide all the lightning, lightning data module, and the error stack trace

Pytorch Custom Dataset

class customDataset(Dataset):
    def __init__(self, imageDir, dataDf, transform=None):
        self.transform = transform
        self.imageDir = imageDir
        self.dataDf = dataDf
        
    def __len__(self):
        return len(self.dataDf)
    
    def _convertXyxy(self, box):
        x, y, w, h = box
        xMin = x
        yMin = y
        xMax = x + w
        yMax = y + h
        return [xMin, yMin, xMax, yMax]
        
    def __getitem__(self, idx):
        fileName = self.dataDf.iloc[idx, 0]
        imgPath = os.path.join(self.imageDir, fileName)
        img = cv2.imread(imgPath)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        label = self.dataDf.iloc[idx, 2]
        boxes = self._convertXyxy(self.dataDf.iloc[idx, 1])
                
        if self.transform:
            img = self.transform(img)
        
        target = {}
        target['boxes'] = torch.tensor([boxes], dtype=torch.float32)
        target['labels'] = torch.tensor([label], dtype=torch.int64)
        return img, target

Lightning Data Module

class dataModule(pl.LightningDataModule):
    def __init__(self, trainingDir = trainDir, validationDir = valDir, trainingDf = subTrainDf, validationDf = subValDf, batchSize = 4):
        super().__init__()
        self.trainingDir = trainingDir
        self.validationDir = validationDir
        self.trainingDf = trainingDf
        self.validationDf = valDf
        self.batchSize = batchSize
        self.classes = ['Background', 'Person']
    
    def prepare_data(self):
        pass
        
    def setup(self, stage):
        mean = [.485, .456, .406]
        std = [.229, .224, .225]
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std)
        ])
        self.trainSet = customDataset(self.trainingDir, self.trainingDf, transform=transform)
        self.valSet = customDataset(self.validationDir, self.validationDf, transform=transform)
    
    def train_dataloader(self):
        return DataLoader(self.trainSet, batch_size=self.batchSize, shuffle=True, num_workers=8, drop_last=True, collate_fn=collateFn)
    
    def val_dataloader(self):
        return DataLoader(self.valSet, batch_size=self.batchSize, shuffle=False, num_workers=8, drop_last=True, collate_fn=collateFn)
    
def collateFn(batch):
    return tuple(zip(*batch))

Lightning Module

class ResNet50_FRCNN(pl.LightningModule):
    def __init__(self, lr, decay=0):
        super().__init__()
        backbone = models.detection.fasterrcnn_resnet50_fpn(weights='DEFAULT', weights_backbone='DEFAULT')
        in_features = backbone.roi_heads.box_predictor.cls_score.in_features
        self.model = backbone
        self.model.roi_heads.box_predictor = models.detection.faster_rcnn.FastRCNNPredictor(in_features, num_classes=2)
        self.mAP = detectionMetrics.mean_ap.MeanAveragePrecision(box_format='xyxy', iou_type='bbox')
        self.met1 = detectionMetrics.iou.IntersectionOverUnion(box_format='xyxy').to(self.device)
        self.lr = lr
        self.decay = decay
        self.mean = [.485, .456, .406]
        self.std = [.229, .224, .225]
        self.unnormalize = transforms.Normalize(mean=[-m/s for m, s in zip(self.mean, self.std)], std=[1/s for s in self.std])
    def forward(self, x):
        x = self.model(x)
        return x
    def training_step(self, batch, batch_idx):
        images, targets = batch
        outputs = self.model(images, targets)
        losses = sum(loss for loss in outputs.values())
        self.log_dict({'training_loss' : losses}, on_step=False, on_epoch=True, prog_bar=True, batch_size=4)
        return losses
    
    def on_validation_start(self):
        self.met1 = self.met1.to(self.device)
    
    def validation_step(self, batch, batch_idx):
        images, targets = batch
        outputs = self.model(images)
        intersect = self.met1(outputs, targets)
        MAP = self.mAP.update(outputs, targets)
        self.log_dict({'IoU' : intersect['iou']}, on_step=False, on_epoch=True, prog_bar=True, batch_size=4)
        
    def on_validation_end(self):
        self.print(self.mAP.compute())
    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.decay)

PL Trainer

dataMod = dataModule()
model = ResNet50_FRCNN(0.0003)
trainer = pl.Trainer(accelerator='gpu', devices=[0], max_epochs=50, precision='16-mixed', limit_val_batches=50)
trainer.fit(model=model, datamodule=dataMod)

Error Stack Trace

RuntimeError                              Traceback (most recent call last)
Cell In [25], line 4
      2 model = ResNet50_FRCNN(0.0003)
      3 trainer = pl.Trainer(accelerator='gpu', devices=[0], max_epochs=50, precision='16-mixed', limit_val_batches=50)
----> 4 trainer.fit(model=model, datamodule=dataMod)

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py:529, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    527 model = _maybe_unwrap_optimized(model)
    528 self.strategy._lightning_module = model
--> 529 call._call_and_handle_interrupt(
    530     self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    531 )

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/call.py:42, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     40     if trainer.strategy.launcher is not None:
     41         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
---> 42     return trainer_fn(*args, **kwargs)
     44 except _TunerExitException:
     45     _call_teardown_hook(trainer)

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py:568, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    558 self._data_connector.attach_data(
    559     model, train_dataloaders=train_dataloaders, val_dataloaders=val_dataloaders, datamodule=datamodule
    560 )
    562 ckpt_path = self._checkpoint_connector._select_ckpt_path(
    563     self.state.fn,
    564     ckpt_path,
    565     model_provided=True,
    566     model_connected=self.lightning_module is not None,
    567 )
--> 568 self._run(model, ckpt_path=ckpt_path)
    570 assert self.state.stopped
    571 self.training = False

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py:973, in Trainer._run(self, model, ckpt_path)
    968 self._signal_connector.register_signal_handlers()
    970 # ----------------------------
    971 # RUN THE TRAINER
    972 # ----------------------------
--> 973 results = self._run_stage()
    975 # ----------------------------
    976 # POST-Training CLEAN UP
    977 # ----------------------------
    978 log.debug(f"{self.__class__.__name__}: trainer tearing down")

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/trainer.py:1016, in Trainer._run_stage(self)
   1014         self._run_sanity_check()
   1015     with torch.autograd.set_detect_anomaly(self._detect_anomaly):
-> 1016         self.fit_loop.run()
   1017     return None
   1018 raise RuntimeError(f"Unexpected state {self.state}")

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/loops/fit_loop.py:201, in _FitLoop.run(self)
    199 try:
    200     self.on_advance_start()
--> 201     self.advance()
    202     self.on_advance_end()
    203     self._restarting = False

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/loops/fit_loop.py:354, in _FitLoop.advance(self)
    352 self._data_fetcher.setup(combined_loader)
    353 with self.trainer.profiler.profile("run_training_epoch"):
--> 354     self.epoch_loop.run(self._data_fetcher)

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/loops/training_epoch_loop.py:134, in _TrainingEpochLoop.run(self, data_fetcher)
    132 try:
    133     self.advance(data_fetcher)
--> 134     self.on_advance_end()
    135     self._restarting = False
    136 except StopIteration:

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/loops/training_epoch_loop.py:248, in _TrainingEpochLoop.on_advance_end(self)
    246 if should_check_val:
    247     self.trainer.validating = True
--> 248     self.val_loop.run()
    249     self.trainer.training = True
    251 # update plateau LR scheduler after metrics are logged

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/loops/utilities.py:177, in _no_grad_context.<locals>._decorator(self, *args, **kwargs)
    175     context_manager = torch.no_grad
    176 with context_manager():
--> 177     return loop_run(self, *args, **kwargs)

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/loops/evaluation_loop.py:115, in _EvaluationLoop.run(self)
    113     previous_dataloader_idx = dataloader_idx
    114     # run step hooks
--> 115     self._evaluation_step(batch, batch_idx, dataloader_idx)
    116 except StopIteration:
    117     # this needs to wrap the `*_step` call too (not just `next`) for `dataloader_iter` support
    118     break

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/loops/evaluation_loop.py:375, in _EvaluationLoop._evaluation_step(self, batch, batch_idx, dataloader_idx)
    372 self.batch_progress.increment_started()
    374 hook_name = "test_step" if trainer.testing else "validation_step"
--> 375 output = call._call_strategy_hook(trainer, hook_name, *step_kwargs.values())
    377 self.batch_progress.increment_processed()
    379 hook_name = "on_test_batch_end" if trainer.testing else "on_validation_batch_end"

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/trainer/call.py:291, in _call_strategy_hook(trainer, hook_name, *args, **kwargs)
    288     return None
    290 with trainer.profiler.profile(f"[Strategy]{trainer.strategy.__class__.__name__}.{hook_name}"):
--> 291     output = fn(*args, **kwargs)
    293 # restore current_fx when nested context
    294 pl_module._current_fx_name = prev_fx_name

File /usr/local/lib/python3.9/dist-packages/pytorch_lightning/strategies/strategy.py:379, in Strategy.validation_step(self, *args, **kwargs)
    377 with self.precision_plugin.val_step_context():
    378     assert isinstance(self.model, ValidationStep)
--> 379     return self.model.validation_step(*args, **kwargs)

Cell In [24], line 40, in ResNet50_FRCNN.validation_step(self, batch, batch_idx)
     38 images, targets = batch
     39 outputs = self.model(images)
---> 40 intersect = self.met1(outputs, targets)
     41 MAP = self.mAP.update(outputs, targets)
     42 self.log_dict({'IoU' : intersect['iou']}, on_step=False, on_epoch=True, prog_bar=True, batch_size=4)

File /usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File /usr/local/lib/python3.9/dist-packages/torchmetrics/metric.py:288, in Metric.forward(self, *args, **kwargs)
    282     raise TorchMetricsUserError(
    283         "The Metric shouldn't be synced when performing ``forward``. "
    284         "HINT: Did you forget to call ``unsync`` ?."
    285     )
    287 if self.full_state_update or self.full_state_update is None or self.dist_sync_on_step:
--> 288     self._forward_cache = self._forward_full_state_update(*args, **kwargs)
    289 else:
    290     self._forward_cache = self._forward_reduce_state_update(*args, **kwargs)

File /usr/local/lib/python3.9/dist-packages/torchmetrics/metric.py:319, in Metric._forward_full_state_update(self, *args, **kwargs)
    317 self.reset()
    318 self.update(*args, **kwargs)
--> 319 batch_val = self.compute()
    321 # restore context
    322 for attr, val in cache.items():

File /usr/local/lib/python3.9/dist-packages/torchmetrics/metric.py:593, in Metric._wrap_compute.<locals>.wrapped_func(*args, **kwargs)
    585 # compute relies on the sync context manager to gather the states across processes and apply reduction
    586 # if synchronization happened, the current rank accumulated states will be restored to keep
    587 # accumulation going if ``should_unsync=True``,
    588 with self.sync_context(
    589     dist_sync_fn=self.dist_sync_fn,
    590     should_sync=self._to_sync,
    591     should_unsync=self._should_unsync,
    592 ):
--> 593     value = _squeeze_if_scalar(compute(*args, **kwargs))
    595 if self.compute_with_cache:
    596     self._computed = value

File /usr/local/lib/python3.9/dist-packages/torchmetrics/detection/iou.py:226, in IntersectionOverUnion.compute(self)
    224 def compute(self) -> dict:
    225     """Computes IoU based on inputs passed in to ``update`` previously."""
--> 226     aggregated_iou = dim_zero_cat(
    227         [self._iou_compute_fn(iou, bool(lbl_eq)) for iou, lbl_eq in zip(self.results, self.labels_eq)]
    228     )
    229     results: Dict[str, Tensor] = {f"{self._iou_type}": aggregated_iou.mean()}
    231     if self.class_metrics:

File /usr/local/lib/python3.9/dist-packages/torchmetrics/utilities/data.py:35, in dim_zero_cat(x)
     33 if not x:  # empty list
     34     raise ValueError("No samples to concatenate")
---> 35 return torch.cat(x, dim=0)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument tensors in method wrapper_cat)

Strange things (Focus on the progress bar and batch, it runs perfectly fine in the first few batch)