File size: 656 Bytes
c865888 |
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 |
from pynvml import *
"""
To track memory allocation, let's take advantage of the nvidia-ml-py3 package and GPU memory allocation from python.
ref: https://huggingface.co/docs/transformers/v4.20.1/en/perf_train_gpu_one
"""
def print_gpu_initialization():
nvmlInit()
handle = nvmlDeviceGetHandleByIndex(0)
info = nvmlDeviceGetMemoryInfo(handle)
print(f"GPU memory occupied: {info.used//1024**2} MB.")
return info.used // 1024**2
def print_summary(result):
print(f"Time: {result.metrics['train_runtime']:.2f}")
print(f"Samples/second: {result.metrics['train_samples_per_second']:.2f}")
print_gpu_utilization()
|