get total, used and free memory (in MBs) for gpu `id`. if `id` is not passed, currently selected torch device is used
(self, _id=None)
| 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): |