gpu_callback.py

import tensorflow as tf
import nvidia_smi

class GPUStats(tf.keras.callbacks.Callback):
    """
    Usage:
        model.fit(..., callbacks=[GPUStats()])

    requires nvidia_smi package:
        conda: conda install -c fastai nvidia-ml-py3
        pip  : pip install nvidia-ml-py3
    """
    def __init__(self, idx=0):
        nvidia_smi.nvmlInit()
        self.handle = nvidia_smi.nvmlDeviceGetHandleByIndex(idx)

    def on_x_end(self, x, logs=None):
        self.mem = nvidia_smi.nvmlDeviceGetMemoryInfo(self.handle)
        self.res = nvidia_smi.nvmlDeviceGetUtilizationRates(self.handle)
        logs["GPU-Usage [%]"] = self.res.gpu
        logs["GPU-vRAM [%]"] = 100 * self.mem.used / self.mem.total
        logs["GPU-vRAM [MiB]"] = self.mem.used / (1024**2)

    on_batch_end = on_x_end