From 49c14d6bf5df9acf4a140584b14c761dc94c083c Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 1 Jul 2022 15:19:57 +0200 Subject: [PATCH 1/2] TRT `--half` fix autocast images to FP16 Resolves bug raised in https://github.com/ultralytics/yolov5/issues/7822 --- models/common.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/models/common.py b/models/common.py index 7690f714def8..64e825b95bb9 100644 --- a/models/common.py +++ b/models/common.py @@ -441,6 +441,9 @@ def wrap_frozen_graph(gd, inputs, outputs): def forward(self, im, augment=False, visualize=False, val=False): # YOLOv5 MultiBackend inference b, ch, h, w = im.shape # batch, channel, height, width + if self.fp16: + im = im.half() # to FP16 + if self.pt: # PyTorch y = self.model(im, augment=augment, visualize=visualize)[0] elif self.jit: # TorchScript From 940b573bf981dfd74a5795c0bb5cba9478a9a1c8 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 1 Jul 2022 15:35:17 +0200 Subject: [PATCH 2/2] Update common.py --- models/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/common.py b/models/common.py index 64e825b95bb9..a6488dd85648 100644 --- a/models/common.py +++ b/models/common.py @@ -441,7 +441,7 @@ def wrap_frozen_graph(gd, inputs, outputs): def forward(self, im, augment=False, visualize=False, val=False): # YOLOv5 MultiBackend inference b, ch, h, w = im.shape # batch, channel, height, width - if self.fp16: + if self.fp16 and im.dtype != torch.float16: im = im.half() # to FP16 if self.pt: # PyTorch