Gpu memory -> Gpu info: memory + utilization

Author: vfdev-5

examples/contrib/mnist/mnist_with_tqdm_logger.py

@@ -44,7 +44,7 @@ def get_data_loaders(train_batch_size, val_batch_size):
     return train_loader, val_loader
 
 
-def run(train_batch_size, val_batch_size, epochs, lr, momentum):
+def run(train_batch_size, val_batch_size, epochs, lr, momentum, display_gpu_info):
     train_loader, val_loader = get_data_loaders(train_batch_size, val_batch_size)
     model = Net()
     device = 'cpu'
@@ -61,8 +61,12 @@ def run(train_batch_size, val_batch_size, epochs, lr, momentum):
 
     RunningAverage(output_transform=lambda x: x).attach(trainer, 'loss')
 
+    if display_gpu_info:
+        from ignite.contrib.metrics import GpuInfo
+        GpuInfo().attach(trainer, name='gpu')
+
     pbar = ProgressBar(persist=True)
-    pbar.attach(trainer, ['loss'])
+    pbar.attach(trainer, metric_names=['loss', 'gpu:0 memory', 'gpu:0 util'])
 
     @trainer.on(Events.EPOCH_COMPLETED)
     def log_training_results(engine):
@@ -102,7 +106,9 @@ def log_validation_results(engine):
                         help='learning rate (default: 0.01)')
     parser.add_argument('--momentum', type=float, default=0.5,
                         help='SGD momentum (default: 0.5)')
+    parser.add_argument('--display_gpu_info', action='store_true', 
+                        help='Display gpu usage info. This needs python 3.X and pynvml package')
 
     args = parser.parse_args()
 
-    run(args.batch_size, args.val_batch_size, args.epochs, args.lr, args.momentum)
+    run(args.batch_size, args.val_batch_size, args.epochs, args.lr, args.momentum, args.display_gpu_info)

ignite/contrib/metrics/__init__.py

@@ -1,4 +1,4 @@
 from ignite.contrib.metrics.average_precision import AveragePrecision
 from ignite.contrib.metrics.roc_auc import ROC_AUC
 import ignite.contrib.metrics.regression
-from ignite.contrib.metrics.gpu_memory import GpuMemory
+from ignite.contrib.metrics.gpu_info import GpuInfo

ignite/contrib/metrics/gpu_info.py

@@ -0,0 +1,87 @@
+# -*- coding: utf-8 -*-
+import warnings
+
+import torch
+
+from ignite.metrics import Metric
+from ignite.engine import Events
+
+
+class GpuInfo(Metric):
+    """GPU information: a) used / max memory, b) gpu utilization values as Metric.
+
+    Examples:
+
+        .. code-block:: python
+
+            # Default GPU measurement
+            GpuInfo().attach(trainer)  # default metric names are 'gpu info:N memory', 'gpu info:N util'
+            ProgressBar(persist=True).attach(trainer, metric_names=['gpu info:0 memory', 'gpu info:0 util'])
+
+            # Progress bar will looks like
+            # Epoch [2/50]: [64/128]  50%|█████      , gpu memory=1120 / 11176 MiB [06:17<12:34]
+
+    """
+
+    def __init__(self):
+        try:
+            import pynvml
+        except ImportError:
+            raise RuntimeError("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")
+
+        from pynvml.smi import nvidia_smi
+        # Let it fail if no libnvidia drivers or NMVL library found
+        self.nvsmi = nvidia_smi.getInstance()
+        super(GpuInfo, self).__init__()
+
+    def reset(self):
+        pass
+
+    def update(self, output):
+        pass
+
+    def compute(self):
+        data = 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, name):
+        data = self.compute()
+        if len(data) < 1:
+            warnings.warn("No GPU information available")
+            return
+
+        for i, data_by_rank in enumerate(data):
+            mem_name = "{}:{} memory".format(name, i)
+
+            if 'fb_memory_usage' not in data_by_rank:
+                warnings.warn("No GPU memory usage information available in {}".format(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 "
+                              "memory consumption information in {}".format(mem_report))
+                continue
+
+            engine.state.metrics[mem_name] = "{}/{} MiB".format(int(mem_report['used']), int(mem_report['total']))
+
+            util_name = "{}:{} util".format(name, i)
+            if 'utilization' not in data_by_rank:
+                warnings.warn("No GPU utilization information available in {}".format(data_by_rank))
+                continue
+            util_report = data_by_rank['utilization']
+            if not ('gpu_util' in util_report):
+                warnings.warn("GPU utilization information does not provide 'gpu_util' information in "
+                              "{}".format(util_report))
+                continue
+
+            engine.state.metrics[util_name] = "{:02d}%".format(int(util_report['gpu_util']))
+
+    def attach(self, engine, name="gpu info", event_name=Events.ITERATION_COMPLETED):
+        engine.add_event_handler(event_name, self.completed, name)

ignite/contrib/metrics/gpu_memory.py

@@ -3,7 +3,7 @@
 import torch
 
 from ignite.engine import Engine, State
-from ignite.contrib.metrics import GpuMemory
+from ignite.contrib.metrics import GpuInfo
 
 import pytest
 
@@ -26,38 +26,47 @@ def no_site_packages():
 def test_no_pynvml_package(no_site_packages):
 
     with pytest.raises(RuntimeError, match="This contrib module requires pynvml to be installed."):
-        GpuMemory()
+        GpuInfo()
 
 
 @pytest.mark.skipif(sys.version[0] == "2" or torch.cuda.is_available(), reason="No pynvml for python 2.7")
 def test_no_gpu():
 
     with pytest.raises(RuntimeError, match="This contrib module requires available GPU"):
-        GpuMemory()
+        GpuInfo()
 
 
 @pytest.mark.skipif(sys.version[0] == "2" or not (torch.cuda.is_available()),
                     reason="No pynvml for python 2.7 and no GPU")
 def test_gpu_mem_consumption():
 
-    gpu_mem = GpuMemory()
+    gpu_info = GpuInfo()
 
     t = torch.rand(4, 10, 100, 100)
-    data = gpu_mem.compute()
+    data = gpu_info.compute()
     assert len(data) > 0
     assert "fb_memory_usage" in data[0]
-    report = data[0]['fb_memory_usage']
-    assert 'used' in report and 'total' in report
-    assert report['total'] > 0.0
-    assert report['used'] > t.shape[0] * t.shape[1] * t.shape[2] * t.shape[3] / 1024.0 / 1024.0
+    mem_report = data[0]['fb_memory_usage']
+    assert 'used' in mem_report and 'total' in mem_report
+    assert mem_report['total'] > 0.0
+    assert mem_report['used'] > t.shape[0] * t.shape[1] * t.shape[2] * t.shape[3] / 1024.0 / 1024.0
+
+    assert "utilization" in data[0]
+    util_report = data[0]['utilization']
+    assert 'gpu_util' in util_report
 
     # with Engine
     engine = Engine(lambda engine, batch: 0.0)
     engine.state = State(metrics={})
 
-    gpu_mem.completed(engine, name='gpu mem', local_rank=0)
+    gpu_info.completed(engine, name='gpu info')
+
+    assert 'gpu info:0 memory' in engine.state.metrics
+    assert 'gpu info:0 util' in engine.state.metrics
+
+    assert isinstance(engine.state.metrics['gpu info:0 memory'], str)
+    assert "{}".format(int(mem_report['used'])) in engine.state.metrics['gpu info:0 memory']
+    assert "{}".format(int(mem_report['total'])) in engine.state.metrics['gpu info:0 memory']
 
-    assert 'gpu mem' in engine.state.metrics
-    assert isinstance(engine.state.metrics['gpu mem'], str)
-    assert "{}".format(int(report['used'])) in engine.state.metrics['gpu mem']
-    assert "{}".format(int(report['total'])) in engine.state.metrics['gpu mem']
+    assert isinstance(engine.state.metrics['gpu info:0 util'], str)
+    assert "{}".format(int(util_report['gpu_util'])) in engine.state.metrics['gpu info:0 util']