Fix different devices bug when moving model from GPU to CPU (ultralyt…

…ics#5110) * fix different devices bug * extend _apply() instead of to() for a general fix * Only apply if Detect() is last layer Co-authored-by: Jebastin Nadar <[email protected]> * Indent fix * Add comment to yolo.py * Add comment to common.py Co-authored-by: Glenn Jocher <[email protected]>
fhkiel-mlaip · Oct 10, 2021 · b2dc5a2 · b2dc5a2
1 parent 8fc4641
commit b2dc5a2
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/models/common.py b/models/common.py
@@ -289,6 +289,14 @@ def autoshape(self):
         LOGGER.info('AutoShape already enabled, skipping... ')  # model already converted to model.autoshape()
         return self
 
+    def _apply(self, fn):
+        # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
+        self = super()._apply(fn)
+        m = self.model.model[-1]  # Detect()
+        m.stride = fn(m.stride)
+        m.grid = list(map(fn, m.grid))
+        return self
+
     @torch.no_grad()
     def forward(self, imgs, size=640, augment=False, profile=False):
         # Inference from various sources. For height=640, width=1280, RGB images example inputs are:

diff --git a/models/yolo.py b/models/yolo.py
@@ -232,6 +232,15 @@ def autoshape(self):  # add AutoShape module
     def info(self, verbose=False, img_size=640):  # print model information
         model_info(self, verbose, img_size)
 
+    def _apply(self, fn):
+        # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
+        self = super()._apply(fn)
+        m = self.model[-1]  # Detect()
+        if isinstance(m, Detect):
+            m.stride = fn(m.stride)
+            m.grid = list(map(fn, m.grid))
+        return self
+
 
 def parse_model(d, ch):  # model_dict, input_channels(3)
     LOGGER.info('\n%3s%18s%3s%10s  %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments'))