Uso real de la memoria de video durante el entrenamiento del modelo usando monitoreo multiproceso nvidia-smi

        Cuando se usa torch.cuda.max_allocated_memory, solo se puede devolver el tamaño de la memoria ocupada por el tensor de la antorcha. Es necesario agregar el tamaño del modelo y el tamaño de carga del modelo no es simplemente el tamaño de la carpeta del modelo. Por el momento, solo observe el comando -n 0.1 nvidia-smi para monitoreo humano.

        Además, los tiempos de ejecución como TRT no admiten los métodos torch.cuda.

        Aquí abrí un hilo para monitorear la memoria.

import nvidia_smi
from threading import Thread
from multiprocessing import Process
import time

class Monitor(Thread):
    def __init__(self, delay, index):
        super(Monitor, self).__init__()
        self.stopped = False
        self.index=index
        self.delay = delay # Time between calls to nvidia-smi
        # st mem
        nvidia_smi.nvmlInit()
        handle = nvidia_smi.nvmlDeviceGetHandleByIndex(index)
        info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
        self.st_mem=info.used
        print( "start used memory is {} GiB".format(info.used* 1.0 / 1024**3))
        nvidia_smi.nvmlShutdown()
        #st_end
        self.max_mem=0
        time.sleep(self.delay)
        
        self.start()

    def run(self):
        while not self.stopped:
            nvidia_smi.nvmlInit()
            handle = nvidia_smi.nvmlDeviceGetHandleByIndex(self.index)
            info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
            self.max_mem=max(self.max_mem, info.used)
            nvidia_smi.nvmlShutdown()
            time.sleep(self.delay)

    def stop(self):
        self.stopped = True
        res=(self.max_mem-self.st_mem)* 1.0 / 1024**3
        print( "total used memory is {} GiB".format(res))
        return res
        
if __name__=="__main__":
    # Instantiate monitor with a 10-second delay between updates
    monitor = Monitor(0.1,0)

    # Train, etc.
    time.sleep(1)

    # Close monitor
    monitor.stop()