| 48 | |
| 49 | |
| 50 | class CudaUse(object): |
| 51 | def __init__(self): |
| 52 | self.cuda_available = th.cuda.is_available() |
| 53 | if self.cuda_available: |
| 54 | from fastai.utils.pynvml_gate import load_pynvml_env |
| 55 | self.pynvml = load_pynvml_env() |
| 56 | |
| 57 | def get_cuda_id(self): |
| 58 | if self.cuda_available: |
| 59 | gpu_mem = sorted(self.gpu_mem_get_all(), key=lambda item: item.free, reverse=True) |
| 60 | low_use_id = gpu_mem[0].id |
| 61 | return th.device(f'cuda:{low_use_id}') |
| 62 | else: |
| 63 | return th.device('cpu') |
| 64 | |
| 65 | def gpu_mem_get_all(self): |
| 66 | "get total, used and free memory (in MBs) for each available gpu" |
| 67 | return list(map(self.gpu_mem_get, range(self.pynvml.nvmlDeviceGetCount()))) |
| 68 | |
| 69 | def gpu_mem_get(self, _id=None): |
| 70 | """get total, used and free memory (in MBs) for gpu `id`. if `id` is not passed, |
| 71 | currently selected torch device is used""" |
| 72 | from collections import namedtuple |
| 73 | GPUMemory = namedtuple('GPUMemory', ['total', 'free', 'used', 'id']) |
| 74 | |
| 75 | if _id is None: |
| 76 | _id = th.cuda.current_device() |
| 77 | try: |
| 78 | handle = self.pynvml.nvmlDeviceGetHandleByIndex(_id) |
| 79 | info = self.pynvml.nvmlDeviceGetMemoryInfo(handle) |
| 80 | # return GPUMemory(*(map(b2mb, [info.total, info.free, info.used])), id=_id) |
| 81 | return b2mb(info.used) |
| 82 | except: |
| 83 | return GPUMemory(0, 0, 0, -1) |
| 84 | |
| 85 | |
| 86 | def read_file(path, mode='r', encoding=None): |