当前位置：首页 > news >正文

huggingface 笔记：查看GPU占用情况

news 2025/8/3 0:30:52

0 准备部分

0.1 创建虚拟数据

import numpy as npfrom datasets import Datasetseq_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 AutoModelForSequenceClassificationmodel = 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, logginglogging.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)