Ankle/yolov5_birdproject
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Added Windows cmd to count GPU devices (#7891)

Browse Source 
* Added Windows cmd to count GPU devices

* Cleanup

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
xylieongandGlenn Jocher authored and GitHub committed 2022-05-19 15:48:44 +02:00
1 parent fe1b503cbb
commit f43cd53d44
1 file changed
+3 -3
Download Patch File Download Diff File Expand all files Collapse all files
utils/torch_utils.py
+3 -3
View File
@@ -40,10 +40,10 @@ def torch_distributed_zero_first(local_rank: int):
def device_count(): def device_count():
# Returns number of CUDA devices available. Safe version of torch.cuda.device_count(). Only works on Linux. # Returns number of CUDA devices available. Safe version of torch.cuda.device_count(). Supports Linux and Windows
assert platform.system() == 'Linux', 'device_count() function only works on Linux' assert platform.system() in ('Linux', 'Windows'), 'device_count() only supported on Linux or Windows'
try: try:
cmd = 'nvidia-smi -L | wc -l' cmd = 'nvidia-smi -L | wc -l' if platform.system() == 'Linux' else 'nvidia-smi -L | find /c /v ""' # Windows
return int(subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]) return int(subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1])
except Exception: except Exception:
return 0 return 0