Repository: pytorch/ignite

Issue:
Add GPU memory consumption as metric
Idea is to provide used GPU max memory during the engine run.
For example,
```python
from pynvml.smi import nvidia_smi
from ignite.engine import Engine, Events
# from ignite.contrib.handlers import ProgressBar
# Here I use a custom ProgressBar to be able to log string value

import time


nvsmi = nvidia_smi.getInstance()
local_rank = 0

def update(engine, batch):
    time.sleep(0.1)
    a = torch.rand(10, 10, 10, 10).cuda()


trainer = Engine(update)


@trainer.on(Events.ITERATION_COMPLETED)
def log_gpu_mem_usage_as_metric(engine):
    
    report = nvsmi.DeviceQuery('memory.free, memory.used, memory.total')['gpu'][local_rank]['fb_memory_usage']    
    engine.state.metrics['gpu memory'] = "{} / {} MiB".format(int(report['used']), int(report['total']))

ProgressBar(persist=True).attach(trainer, metric_names=['gpu memory', ])

data = list(range(40))
trainer.run(data)
```
with output
```
Epoch
[40/40] 100%|██████████, gpu memory=493 / 11176 MiB [00:04<00:00]
```

However wrapping the handler into a class as `RunningAverage`



Discussion:

Resolving patch:
diff --git a/examples/contrib/mnist/mnist_with_tqdm_logger.py b/examples/contrib/mnist/mnist_with_tqdm_logger.py
--- a/examples/contrib/mnist/mnist_with_tqdm_logger.py
+++ b/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)
diff --git a/ignite/contrib/metrics/__init__.py b/ignite/contrib/metrics/__init__.py
--- a/ignite/contrib/metrics/__init__.py
+++ b/ignite/contrib/metrics/__init__.py
@@ -1,3 +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_info import GpuInfo
diff --git a/ignite/contrib/metrics/gpu_info.py b/ignite/contrib/metrics/gpu_info.py
new file mode 100644
--- /dev/null
+++ b/ignite/contrib/metrics/gpu_info.py
@@ -0,0 +1,88 @@
+# -*- 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, name='gpu')  # metric names are 'gpu:X memory', 'gpu:X util'
+            ProgressBar(persist=True).attach(trainer, metric_names=['gpu:0 memory', 'gpu:0 util'])
+
+            # Progress bar will looks like
+            # Epoch [2/10]: [12/24]  50%|█████      , gpu:0 memory=1120/11176 MiB, gpu:0 util=09% [00:17<1:23]
+
+    """
+
+    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']))
+
+        for i, data_by_rank in enumerate(data):
+            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)


Test patch:
diff --git a/tests/ignite/contrib/metrics/test_gpu_info.py b/tests/ignite/contrib/metrics/test_gpu_info.py
new file mode 100644
--- /dev/null
+++ b/tests/ignite/contrib/metrics/test_gpu_info.py
@@ -0,0 +1,163 @@
+import sys
+
+import torch
+
+from ignite.engine import Engine, State
+from ignite.contrib.metrics import GpuInfo
+
+import pytest
+from mock import Mock, patch
+
+
+@pytest.fixture
+def no_site_packages():
+    import pynvml
+    import sys
+    assert 'pynvml' in sys.modules
+    pynvml_module = sys.modules['pynvml']
+    del sys.modules['pynvml']
+    prev_path = list(sys.path)
+    sys.path = [p for p in sys.path if "site-packages" not in p]
+    yield "no_site_packages"
+    sys.path = prev_path
+    sys.modules['pynvml'] = pynvml_module
+
+
+@pytest.mark.skipif(sys.version[0] == "2", reason="No pynvml for python 2.7")
+def test_no_pynvml_package(no_site_packages):
+
+    with pytest.raises(RuntimeError, match="This contrib module requires pynvml to be installed."):
+        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"):
+        GpuInfo()
+
+
+def _test_gpu_info(device='cpu'):
+    gpu_info = GpuInfo()
+
+    # increase code cov
+    gpu_info.reset()
+    gpu_info.update(None)
+
+    t = torch.rand(4, 10, 100, 100).to(device)
+    data = gpu_info.compute()
+    assert len(data) > 0
+    assert "fb_memory_usage" in data[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_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 isinstance(engine.state.metrics['gpu info:0 util'], str)
+    assert "{}".format(int(util_report['gpu_util'])) in engine.state.metrics['gpu info:0 util']
+
+
+@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_info():
+    _test_gpu_info(device='cuda')
+
+
+@pytest.fixture
+def mock_pynvml_module():
+
+    with patch.dict('sys.modules', {
+        'pynvml': Mock(name='pynvml'),
+        'pynvml.smi': Mock(name='pynvml.smi'),
+        'pynvml.smi.nvidia_smi': Mock(name='pynvml.smi.nvidia_smi'),
+    }):
+        import pynvml
+        from pynvml.smi import nvidia_smi
+
+        def query(*args, **kwargs):
+            return {
+                "gpu": [{
+                    "fb_memory_usage": {
+                        "used": 100.0,
+                        "total": 11000.0
+                    },
+                    "utilization": {
+                        "gpu_util": 50.0
+                    }
+                }]
+            }
+
+        def getInstance():
+            nvsmi = Mock()
+            nvsmi.DeviceQuery = Mock(side_effect=query)
+            return nvsmi
+
+        nvidia_smi.getInstance = Mock(side_effect=getInstance)
+        yield pynvml
+
+
+@pytest.fixture
+def mock_gpu_is_available():
+
+    with patch('ignite.contrib.metrics.gpu_info.torch.cuda') as mock_cuda:
+        mock_cuda.is_available.return_value = True
+        yield mock_cuda
+
+
+@pytest.mark.skipif(torch.cuda.is_available(), reason="No need to mock if has GPU")
+def test_gpu_info_mock(mock_pynvml_module, mock_gpu_is_available):
+
+    assert torch.cuda.is_available()
+    _test_gpu_info()
+
+    def _test_with_custom_query(resp, warn_msg, check_compute=False):
+        from pynvml.smi import nvidia_smi
+
+        def query(*args, **kwargs):
+            return resp
+
+        def getInstance():
+            nvsmi = Mock()
+            nvsmi.DeviceQuery = Mock(side_effect=query)
+            return nvsmi
+
+        nvidia_smi.getInstance = Mock(side_effect=getInstance)
+        gpu_info = GpuInfo()
+        if check_compute:
+            with pytest.warns(UserWarning, match=warn_msg):
+                gpu_info.compute()
+
+        # with Engine
+        engine = Engine(lambda engine, batch: 0.0)
+        engine.state = State(metrics={})
+
+        with pytest.warns(UserWarning, match=warn_msg):
+            gpu_info.completed(engine, name='gpu info')
+
+    # No GPU info
+    _test_with_custom_query(resp={}, warn_msg=r"No GPU information available", check_compute=True)
+
+    # No GPU memory info
+    _test_with_custom_query(resp={"gpu": [{"utilization": {}}, ]},
+                            warn_msg=r"No GPU memory usage information available")
+
+    # No GPU utilization info
+    _test_with_custom_query(resp={"gpu": [{"fb_memory_usage": {}}, ]},
+                            warn_msg=r"No GPU utilization information available")