YOLOv5 Export Benchmarks for GPU (#6963)
* Add benchmarks.py GPU support * Updates * Updates * Updates * Updates * Add --half * Add TRT requirements * Cleanup * Add TF to warmup types * Update export.py * Update export.py * Update benchmarks.py
This commit is contained in:
1 parent
99de551f97
commit
932dc78496
3 files changed
+31
-18
No files matched your search
+4
-3
@@ -464,10 +464,11 @@ class DetectMultiBackend(nn.Module):
|
||||
|
||||
def warmup(self, imgsz=(1, 3, 640, 640)):
|
||||
# Warmup model by running inference once
|
||||
if self.pt or self.jit or self.onnx or self.engine: # warmup types
|
||||
if isinstance(self.device, torch.device) and self.device.type != 'cpu': # only warmup GPU models
|
||||
if any((self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb)): # warmup types
|
||||
if self.device.type != 'cpu': # only warmup GPU models
|
||||
im = torch.zeros(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input
|
||||
self.forward(im) # warmup
|
||||
for _ in range(2 if self.jit else 1): #
|
||||
self.forward(im) # warmup
|
||||
|
||||
@staticmethod
|
||||
def model_type(p='path/to/model.pt'):
|
||||
|
||||
Reference in new issue
Block a user