From 4aaf3c451b1704706b61a658d954c4389511d5d0 Mon Sep 17 00:00:00 2001 From: AlexeyAB84 Date: Sat, 30 Jul 2022 09:03:33 +0300 Subject: [PATCH] use Half-precision only for Compute Capability >= 7.0 --- detect.py | 4 +++- test.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/detect.py b/detect.py index 53b63eb..b4a7bca 100644 --- a/detect.py +++ b/detect.py @@ -28,7 +28,9 @@ def detect(save_img=False): # Initialize set_logging() device = select_device(opt.device) - half = device.type != 'cpu' # half precision only supported on CUDA + compute_capability = torch.cuda.get_device_capability(device=device) + print(f"compute_capability = {compute_capability} ") + half = (device.type != 'cpu') and (compute_capability[0] >= 7) # half precision only supported on CUDA # Load model model = attempt_load(weights, map_location=device) # load FP32 model diff --git a/test.py b/test.py index 60154a6..3223083 100644 --- a/test.py +++ b/test.py @@ -62,7 +62,9 @@ def test(data, model = TracedModel(model, device, opt.img_size) # Half - half = device.type != 'cpu' and half_precision # half precision only supported on CUDA + compute_capability = torch.cuda.get_device_capability(device=device) + print(f"compute_capability = {compute_capability} ") + half = device.type != 'cpu' and half_precision and (compute_capability[0] >= 7) # half precision only supported on CUDA if half: model.half()