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")