Repository: airctic/icevision

Issue:
show_results with yolov5 + lightning training throws error
## 🐛 Bug
**Describe the bug**
`show_results` with yolo don't work if the model was training with pytorch lightning

**To Reproduce**
Train a yolo model with pytorch-lightning and try to call `show_results`

**Full stacktrace**
```python
---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

<ipython-input-19-2cf4276b061d> in <module>()
----> 1 model_type.show_results(model, valid_ds, detection_threshold=.5)

9 frames

/usr/local/lib/python3.7/dist-packages/icevision/models/ultralytics/yolov5/show_results.py in show_results(model, dataset, detection_threshold, num_samples, ncols, denormalize_fn, show, device)
     41         show=show,
     42         detection_threshold=detection_threshold,
---> 43         device=device,
     44     )
     45 

/usr/local/lib/python3.7/dist-packages/icevision/models/base_show_results.py in base_show_results(predict_fn, model, dataset, num_samples, ncols, denormalize_fn, show, **predict_kwargs)
     19 ) -> None:
     20     records = random.choices(dataset, k=num_samples)
---> 21     preds = predict_fn(model, records, **predict_kwargs)
     22 
     23     show_preds(

/usr/local/lib/python3.7/dist-packages/icevision/models/ultralytics/yolov5/prediction.py in predict(model, dataset, detection_threshold, keep_images, device)
     48         detection_threshold=detection_threshold,
     49         keep_images=keep_images,
---> 50         device=device,
     51     )
     52 

/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     24         def decorate_context(*args, **kwargs):
     25             with self.__class__():
---> 26                 return func(*args, **kwargs)
     27         return cast(F, decorate_context)
     28 

/usr/local/lib/python3.7/dist-packages/icevision/models/ultralytics/yolov5/prediction.py in _predict_batch(model, batch, records, detection_threshold, keep_images, device)
     24     model = model.eval().to(device)
     25 
---> 26     raw_preds = model(batch)[0]
     27     return convert_raw_predictions(
     28         batch=batch,

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/usr/local/lib/python3.7/dist-packages/yolov5/models/yolo.py in forward(self, x, augment, profile)
    121             return torch.cat(y, 1), None  # augmented inference, train
    122         else:
--> 123             return self.forward_once(x, profile)  # single-scale inference, train
    124 
    125     def forward_once(self, x, profile=False):

/usr/local/lib/python3.7/dist-packages/yolov5/models/yolo.py in forward_once(self, x, profile)
    137                 print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))
    138 
--> 139             x = m(x)  # run
    140             y.append(x if m.i in self.save else None)  # save output
    141 

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/usr/local/lib/python3.7/dist-packages/yolov5/models/yolo.py in forward(self, x)
     52 
     53                 y = x[i].sigmoid()
---> 54                 y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
     55                 y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
     56                 z.append(y.view(bs, -1, self.no))

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```



Discussion:
My guess is that the device casting under `_predict_batch` might be incorrect
Can you please assign this one to me?
> Can you please assign this one to me?

Yes sir! Thank you! =)

Resolving patch:
diff --git a/icevision/models/ultralytics/yolov5/prediction.py b/icevision/models/ultralytics/yolov5/prediction.py
--- a/icevision/models/ultralytics/yolov5/prediction.py
+++ b/icevision/models/ultralytics/yolov5/prediction.py
@@ -19,7 +19,17 @@ def _predict_batch(
     keep_images: bool = False,
     device: Optional[torch.device] = None,
 ) -> List[Prediction]:
-    device = device or model_device(model)
+    # device issue addressed on discord: https://discord.com/channels/735877944085446747/770279401791160400/832361687855923250
+    if device is not None:
+        raise ValueError(
+            "For YOLOv5 device can only be specified during model creation, "
+            "for more info take a look at the discussion here: "
+            "https://discord.com/channels/735877944085446747/770279401791160400/832361687855923250"
+        )
+    grid = model.model[-1].grid[-1]
+    # if `grid.numel() == 1` it means the grid isn't initialized yet and we can't
+    # trust it's device (will always be CPU)
+    device = grid.device if grid.numel() > 1 else model_device(model)
 
     batch = batch[0].to(device)
     model = model.eval().to(device)


Test patch: