Unverified 提交 d94b4705 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

TRT `--half` fix autocast images to FP16 (#8435)

* TRT `--half` fix autocast images to FP16 Resolves bug raised in https://github.com/ultralytics/yolov5/issues/7822 * Update common.py
上级 89833243
...@@ -441,6 +441,9 @@ class DetectMultiBackend(nn.Module): ...@@ -441,6 +441,9 @@ class DetectMultiBackend(nn.Module):
def forward(self, im, augment=False, visualize=False, val=False): def forward(self, im, augment=False, visualize=False, val=False):
# YOLOv5 MultiBackend inference # YOLOv5 MultiBackend inference
b, ch, h, w = im.shape # batch, channel, height, width b, ch, h, w = im.shape # batch, channel, height, width
if self.fp16 and im.dtype != torch.float16:
im = im.half() # to FP16
if self.pt: # PyTorch if self.pt: # PyTorch
y = self.model(im, augment=augment, visualize=visualize)[0] y = self.model(im, augment=augment, visualize=visualize)[0]
elif self.jit: # TorchScript elif self.jit: # TorchScript
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论