File size: 598 Bytes
23dc493 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
# utils/perf_tuning.py
import os, logging
try:
import cv2
except Exception:
cv2 = None
import torch
def apply():
os.environ.setdefault("OMP_NUM_THREADS", "4")
if cv2:
try:
cv2.setNumThreads(4)
except Exception as e:
logging.info("cv2 threads not set: %s", e)
if torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
try:
logging.info("CUDA device %s β cuDNN benchmark ON", torch.cuda.get_device_name(0))
except Exception:
logging.info("CUDA available β cuDNN benchmark ON")
|