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:
1 parent
d40cd0d454
commit
61adf017f2
5 files changed
+8
-8
No files matched your search
+2
-2
@@ -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):
|
||||
|
||||
Reference in new issue
Block a user