mirrors
/
yolov7
mirror of https://github.com/WongKinYiu/yolov7.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

use Half-precision only for Compute Capability >= 7.0

pull/368/head
AlexeyAB84 2022-07-30 09:03:33 +03:00
parent 6bacefff5c
commit 4aaf3c451b
2 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
detect.py
View File

@ -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

4
test.py
View File

@ -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()