Training with PyTortch Lightning isn't working on GPU because anchor boxes are on CPU

[kennyFF92]

Hi, as in the title I'm training to train EfficientDet using PyTorch Lightning framework. The train works without any issues on CPU but when using GPU it doesn't pass the sanity check because it find the anchor boxes on the CPU device while the other tensors are on the GPU device.

The following is the stack trace of the error, in bold you'll find the relevant parts:

---

Traceback (most recent call last):
File "src/train.py", line 23, in main
trainer.fit(model, data)
File "lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 552, in fit
self._run(model)
File "lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 917, in _run
self._dispatch()
File "lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 985, in _dispatch
self.accelerator.start_training(self)
File "lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py", line 92, in start_training
self.training_type_plugin.start_training(trainer)
File "lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 161, in start_training
self._results = trainer.run_stage()
File "lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 995, in run_stage
return self._run_train()
File "lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1030, in _run_train
self._run_sanity_check(self.lightning_module)
File "lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1114, in _run_sanity_check
self._evaluation_loop.run()
File "lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 111, in run
self.advance(*args, **kwargs)
File "lib/python3.7/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 111, in advance
dataloader_iter, self.current_dataloader_idx, dl_max_batches, self.num_dataloaders
File "lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 111, in run
self.advance(*args, **kwargs)
File "lib/python3.7/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 110, in advance
output = self.evaluation_step(batch, batch_idx, dataloader_idx)
File "lib/python3.7/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 154, in evaluation_step
output = self.trainer.accelerator.validation_step(step_kwargs)
File "lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py", line 211, in validation_step
return self.training_type_plugin.validation_step(*step_kwargs.values())
File "lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 387, in validation_step
return self.model(*args, **kwargs)
File "lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "lib/python3.7/site-packages/torch/nn/parallel/distributed.py", line 705, in forward
output = self.module(*inputs[0], **kwargs[0])
File "lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "lib/python3.7/site-packages/pytorch_lightning/overrides/base.py", line 93, in forward
output = self.module.validation_step(*inputs, **kwargs)
File "/home/ec2-user/SageMaker/lightning-efficientdet/src/efficientdet/models/lit_efficientdet.py", line 156, in validation_step
return self._evaluation(batch, stage="val")
File "/home/ec2-user/SageMaker/lightning-efficientdet/src/efficientdet/models/lit_efficientdet.py", line 111, in _evaluation
losses_dict = bench_train(images, annotations)
File "lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "lib/python3.7/site-packages/effdet/bench.py", line 132, in forward
target['bbox'], target['cls'])
File "lib/python3.7/site-packages/effdet/anchors.py", line 389, in batch_label_anchors
cls_targets, box_targets, matches = self.target_assigner.assign(anchor_box_list, gt_box_list, gt_class_i)
File "lib/python3.7/site-packages/effdet/object_detection/target_assigner.py", line 146, in assign
match_quality_matrix = self._similarity_calc.compare(groundtruth_boxes, anchors)
File "lib/python3.7/site-packages/effdet/object_detection/region_similarity_calculator.py", line 101, in compare
return iou(boxlist1, boxlist2)
File "lib/python3.7/site-packages/effdet/object_detection/region_similarity_calculator.py", line 69, in iou
intersections = intersection(boxlist1, boxlist2)
File "lib/python3.7/site-packages/effdet/object_detection/region_similarity_calculator.py", line 50, in intersection
all_pairs_min_ymax = torch.min(y_max1, y_max2.t())
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

The error is caused by the anchor_box_list variable on row 378 of anchors.py. At row 226 I see that the boxes tensor used to create it, is declared using self.register_buffer so that Lightning can put it in the correct device automatically but it doesn't seem to work. At the moment I can't understand if it's a bug of this library or a bug of Lightning or maybe I'm doing something wrong.

[rwightman]

There is a good example notebook by Chris Hughes with a Lightning example, https://gist.github.com/Chris-hughes10/73628b1d8d6fc7d359b3dcbbbb8869d7 ... the LightningModule (EfficientDetModel) there looks reasonable.

[kennyFF92]

Thanks, I was doing some things wrong and the notebook helped me.