huggingface 笔记：查看GPU占用情况

0 准备部分

0.1 创建虚拟数据

import numpy as np

from datasets import Dataset

seq_len, dataset_size = 512, 512
dummy_data = {
    "input_ids": np.random.randint(100, 30000, (dataset_size, seq_len)),
    "labels": np.random.randint(0, 1, (dataset_size)),
}
dummy_data
'''
使用 np.random.randint 函数生成一个形状为 (dataset_size, seq_len)，即 512x512 的数组。
数组中的每个元素是一个随机整数，范围从 100 到 30000。


使用 np.random.randint 函数生成一个形状为 (dataset_size,) 的数组，其中的元素是 0 或 1
表示每个样本的标签
'''
'''
{'input_ids': array([[11687,  1246,  6661, ..., 20173,  3772, 29152],
        [  720, 25945, 11963, ..., 11675, 27842,  3553],
        [22100, 26587, 19452, ...,  1836, 24395, 22849],
        ...,
        [11610, 24425,  1026, ...,  6237, 28503,  2775],
        [10266, 22622, 14079, ..., 24491, 26029, 17796],
        [11500,  7688, 13780, ...,  4839, 13967, 18493]]),
 'labels': array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0])}
'''

ds = Dataset.from_dict(dummy_data)
ds.set_format("pt")
#将它们存储在一个带有PyTorch格式的数据集中
ds
'''
Dataset({
    features: ['input_ids', 'labels'],
    num_rows: 512
})
'''

0.2 辅助函数

为了打印GPU利用率和使用Trainer进行训练运行的摘要统计信息，定义了两个辅助函数

0.2.1 打印GPU内存使用情况

from pynvml import *

def print_gpu_utilization(device_id=1):
    nvmlInit()#初始化 NVML 库
    
    handle = nvmlDeviceGetHandleByIndex(device_id)
    #获取索引为 device_id 的 GPU 设备的句柄。
    
    info = nvmlDeviceGetMemoryInfo(handle)
    '''
    使用前面获取的句柄，查询该 GPU 的内存信息。
    返回的对象包含了 GPU 内存的总量、已使用的量和空闲的量。
    '''
    
    print(f"GPU内存占用：{info.used // 1024 ** 2} MB。")
    '''
    打印 GPU 当前已使用的内存量，单位为 MB。
    这里通过将字节单位的值除以 1024 的平方来转换为 MB

print_gpu_utilization()
#GPU内存占用：270 MB。

0.2.2 打印训练过程的信息

def print_summary(result):
    print(f"时间：{result.metrics['train_runtime']:.2f}")
    '''
    打印训练过程的运行时间
    从 result 对象中的 metrics 字典获取 'train_runtime' 键的值
    格式化为两位小数。
    '''
    
    print(f"样本/秒：{result.metrics['train_samples_per_second']:.2f}")
    '''
    打印训练速度，即每秒处理的样本数
    从 result 对象中的 metrics 字典获取 'train_samples_per_second' 键的值
    格式化为两位小数。
    '''
    
    print_gpu_utilization()
    #调用 print_gpu_utilization() 函数来打印 GPU 的内存使用情况

1 加载模型，查看GPU空间占用情况

print_gpu_utilization(1)
#print_gpu_utilization(1)


from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained("google-bert/bert-large-uncased").to("cuda:1")
print_gpu_utilization(1)
#GPU内存占用：1866 MB。

2 训练模型时的GPU占用情况

default_args = {
    "output_dir": "tmp",
    "eval_strategy": "steps",
    "num_train_epochs": 1,
    "log_level": "error",
    "report_to": "none",
}
'''
output_dir: 指定输出目录，这里设置为 "tmp"。
eval_strategy: 设置评估策略为 "steps"，意味着在训练过程中会按照步骤进行模型评估。
num_train_epochs: 设置训练周期为 1，即整个训练集只会被训练一遍。
log_level: 设置日志级别为 "error"，这样只有错误信息会被记录。
report_to: 设置报告输出目标为 "none"，这表示不将训练进度报告输出到任何外部服务或控制台。
'''

from transformers import TrainingArguments, Trainer, logging

logging.set_verbosity_error()
#让 transformers 库只输出错误级别的日志。


training_args = TrainingArguments(per_device_train_batch_size=1,
                                  **default_args)
'''
TrainingArguments: 创建一个训练参数对象，设置每个设备的训练批量大小为 1
并将前面定义的默认参数集成进来。
'''
trainer = Trainer(model=model, args=training_args, train_dataset=ds)
#Trainer: 初始化训练器，传入模型、训练参数和训练数据集。
result = trainer.train()
#使用 trainer.train() 启动训练过程，并将结果存储在 result 变量中
print_summary(result)