torch.empty() for speed improvements (#9025)

`torch.empty()` for speed improvement

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-08-18 20:12:33 +02:00
1 parent d40cd0d454
commit 61adf017f2
5 files changed
+8 -8

No files matched your search

+2 -2
View File
@@ -531,7 +531,7 @@ class DetectMultiBackend(nn.Module):
# Warmup model by running inference once
warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb
if any(warmup_types) and self.device.type != 'cpu':
im = torch.zeros(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input
im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input
for _ in range(2 if self.jit else 1): #
self.forward(im) # warmup
@@ -600,7 +600,7 @@ class AutoShape(nn.Module):
dt = (Profile(), Profile(), Profile())
with dt[0]:
p = next(self.model.parameters()) if self.pt else torch.zeros(1, device=self.model.device) # param
p = next(self.model.parameters()) if self.pt else torch.empty(1, device=self.model.device) # param
autocast = self.amp and (p.device.type != 'cpu') # Automatic Mixed Precision (AMP) inference
if isinstance(ims, torch.Tensor): # torch
with amp.autocast(autocast):