From 8ee939e8d18fb88fba3bdfaa9252ccd0572df1b3 Mon Sep 17 00:00:00 2001 From: SamFC10 Date: Sat, 9 Oct 2021 13:50:01 +0530 Subject: [PATCH 1/6] fix different devices bug --- models/common.py | 7 +++++++ models/yolo.py | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/models/common.py b/models/common.py index 2acf6281f475..c62a8c04f5c4 100644 --- a/models/common.py +++ b/models/common.py @@ -289,6 +289,13 @@ def autoshape(self): LOGGER.info('AutoShape already enabled, skipping... ') # model already converted to model.autoshape() return self + def to(self, *args, **kwargs): + self = super().to(*args, **kwargs) + m = self.model.model[-1] # Detect() + m.stride = m.stride.to(*args, **kwargs) + m.grid = [grid.to(*args, **kwargs) for grid in 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 index 5087a9f7399d..b86af521b403 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -232,6 +232,13 @@ def autoshape(self): # add AutoShape module def info(self, verbose=False, img_size=640): # print model information model_info(self, verbose, img_size) + def to(self, *args, **kwargs): + self = super().to(*args, **kwargs) + m = self.model[-1] # Detect() + m.stride = m.stride.to(*args, **kwargs) + m.grid = [grid.to(*args, **kwargs) for grid in 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')) From 2243f4873ee85dc3f7a16d1b49883d1e7445ca65 Mon Sep 17 00:00:00 2001 From: SamFC10 Date: Sun, 10 Oct 2021 13:37:28 +0530 Subject: [PATCH 2/6] extend _apply() instead of to() for a general fix --- models/common.py | 8 ++++---- models/yolo.py | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/models/common.py b/models/common.py index c62a8c04f5c4..1ec712966272 100644 --- a/models/common.py +++ b/models/common.py @@ -289,11 +289,11 @@ def autoshape(self): LOGGER.info('AutoShape already enabled, skipping... ') # model already converted to model.autoshape() return self - def to(self, *args, **kwargs): - self = super().to(*args, **kwargs) + def _apply(self, fn): + self = super()._apply(fn) m = self.model.model[-1] # Detect() - m.stride = m.stride.to(*args, **kwargs) - m.grid = [grid.to(*args, **kwargs) for grid in m.grid] + m.stride = fn(m.stride) + m.grid = list(map(fn, m.grid)) return self @torch.no_grad() diff --git a/models/yolo.py b/models/yolo.py index b86af521b403..0c8138c0e31a 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -232,11 +232,11 @@ def autoshape(self): # add AutoShape module def info(self, verbose=False, img_size=640): # print model information model_info(self, verbose, img_size) - def to(self, *args, **kwargs): - self = super().to(*args, **kwargs) + def _apply(self, fn): + self = super()._apply(fn) m = self.model[-1] # Detect() - m.stride = m.stride.to(*args, **kwargs) - m.grid = [grid.to(*args, **kwargs) for grid in m.grid] + m.stride = fn(m.stride) + m.grid = list(map(fn, m.grid)) return self From 1030165b5e31d3ce9b91b1a10fbbfc7c654b7f3f Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 10 Oct 2021 14:07:48 -0700 Subject: [PATCH 3/6] Only apply if Detect() is last layer Co-authored-by: Jebastin Nadar --- models/yolo.py | 1 + 1 file changed, 1 insertion(+) diff --git a/models/yolo.py b/models/yolo.py index 0c8138c0e31a..01a2c88ab93e 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -235,6 +235,7 @@ def info(self, verbose=False, img_size=640): # print model information def _apply(self, fn): 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 From 204b8bb0fb0dfdeebcfc9fab3287f0e91d2e3a67 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 10 Oct 2021 14:08:25 -0700 Subject: [PATCH 4/6] Indent fix --- models/yolo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/yolo.py b/models/yolo.py index 01a2c88ab93e..8754513e6d8a 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -236,8 +236,8 @@ def _apply(self, fn): 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)) + m.stride = fn(m.stride) + m.grid = list(map(fn, m.grid)) return self From d5a8bde9326515626e9cd8a8aea61fd13662dddc Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 10 Oct 2021 14:14:23 -0700 Subject: [PATCH 5/6] Add comment to yolo.py --- models/yolo.py | 1 + 1 file changed, 1 insertion(+) diff --git a/models/yolo.py b/models/yolo.py index 8754513e6d8a..74b266bff00d 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -233,6 +233,7 @@ 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): From 296881f268f6cf4acf2536992057753782995e4d Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 10 Oct 2021 14:16:21 -0700 Subject: [PATCH 6/6] Add comment to common.py --- models/common.py | 1 + 1 file changed, 1 insertion(+) diff --git a/models/common.py b/models/common.py index 1ec712966272..56077c84acf9 100644 --- a/models/common.py +++ b/models/common.py @@ -290,6 +290,7 @@ def autoshape(self): 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)