1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
|
# -*- coding: utf-8 -*-
import warnings
from typing import Any, Dict, List, Tuple, Union
import torch
from ignite.engine import Engine, EventEnum, Events
from ignite.metrics.metric import Metric
class GpuInfo(Metric):
"""Provides GPU information: a) used memory percentage, b) gpu utilization percentage values as Metric
on each iterations.
.. Note ::
In case if gpu utilization reports "N/A" on a given GPU, corresponding metric value is not set.
Examples:
.. code-block:: python
# Default GPU measurements
GpuInfo().attach(trainer, name='gpu') # metric names are 'gpu:X mem(%)', 'gpu:X util(%)'
# Logging with TQDM
ProgressBar(persist=True).attach(trainer, metric_names=['gpu:0 mem(%)', 'gpu:0 util(%)'])
# Progress bar will looks like
# Epoch [2/10]: [12/24] 50%|█████ , gpu:0 mem(%)=79, gpu:0 util(%)=59 [00:17<1:23]
# Logging with Tensorboard
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training", metric_names='all'),
event_name=Events.ITERATION_COMPLETED)
"""
def __init__(self) -> None:
try:
from pynvml.smi import nvidia_smi
except ImportError:
raise ModuleNotFoundError(
"This contrib module requires pynvml to be installed. "
"Please install it with command: \n pip install pynvml"
)
# Let's check available devices
if not torch.cuda.is_available():
raise RuntimeError("This contrib module requires available GPU")
# Let it fail if no libnvidia drivers or NMVL library found
self.nvsmi = nvidia_smi.getInstance()
super(GpuInfo, self).__init__()
def reset(self) -> None:
pass
def update(self, output: Tuple[torch.Tensor, torch.Tensor]) -> None:
pass
def compute(self) -> List[Dict[str, Any]]:
data: Dict[str, List[Dict[str, Any]]] = self.nvsmi.DeviceQuery("memory.used, memory.total, utilization.gpu")
if len(data) == 0 or ("gpu" not in data):
warnings.warn("No GPU information available")
return []
return data["gpu"]
def completed(self, engine: Engine, name: str) -> None:
data = self.compute()
if len(data) < 1:
warnings.warn("No GPU information available")
return
for i, data_by_rank in enumerate(data):
mem_name = f"{name}:{i} mem(%)"
if "fb_memory_usage" not in data_by_rank:
warnings.warn(f"No GPU memory usage information available in {data_by_rank}")
continue
mem_report = data_by_rank["fb_memory_usage"]
if not ("used" in mem_report and "total" in mem_report):
warnings.warn(
"GPU memory usage information does not provide used/total "
f"memory consumption information in {mem_report}"
)
continue
engine.state.metrics[mem_name] = int(mem_report["used"] * 100.0 / mem_report["total"])
for i, data_by_rank in enumerate(data):
util_name = f"{name}:{i} util(%)"
if "utilization" not in data_by_rank:
warnings.warn(f"No GPU utilization information available in {data_by_rank}")
continue
util_report = data_by_rank["utilization"]
if not ("gpu_util" in util_report):
warnings.warn(f"GPU utilization information does not provide 'gpu_util' information in {util_report}")
continue
try:
engine.state.metrics[util_name] = int(util_report["gpu_util"])
except ValueError:
# Do not set GPU utilization information
pass
# TODO: see issue https://github.com/pytorch/ignite/issues/1405
def attach( # type: ignore
self, engine: Engine, name: str = "gpu", event_name: Union[str, EventEnum] = Events.ITERATION_COMPLETED
) -> None:
engine.add_event_handler(event_name, self.completed, name)
|
