显存管理的方法

笔者在使用for while循环无限次调用本地大模型的时候，炸内存。在深度学习模型的训练和推理过程中，显存管理是一个非常重要的方面，特别是在使用大型模型时。显存管理不当可能会导致显存不足，从而导致程序崩溃或性能下降。以下是一些显存管理的方法和技巧，可以帮助你在使用循环运行代码时更好地管理显存。

我运行的代码如下（先提供个示例代码，方便大家之后看后面的知识来理解）

import os
import sys
import json
import torch
from llava.model.builder import load_pretrained_model
from llava.mm_utils import get_model_name_from_path
from llava.eval.run_llava import eval_model# 从命令行参数获取输入
model_path = sys.argv[1]
data_file = sys.argv[2]
output_file = sys.argv[3]# 加载预训练模型
def load_model(model_path):tokenizer, model, image_processor, context_len = load_pretrained_model(model_path=model_path,model_base=None,model_name=get_model_name_from_path(model_path),# load_4bit=True,# load_8bit=True     # 如果不是8bit跑，需要占用30944mib)return tokenizer, model, image_processor, context_len# 读取 JSON 数据文件
with open(data_file, 'r') as f:data = json.load(f)# 打开输出文件
with open(output_file, 'w') as out_f:# 遍历每一张图片的数据for image in data:image_name = image['image_path']pred_classes = image['pred_classes']image_file = f"/home/data/yjgroup/fsy/VG_100K/{image_name}"prompt = ("Generate relation triples for the objects in the image. Relation categories: ""Input image: '{}', objects: {}. Generate triples and confidences.").format(image_name, pred_classes)args = type('Args', (), {"model_path": model_path,"model_base": None,"model_name": get_model_name_from_path(model_path),"query": prompt,"conv_mode": None,"image_file": image_file,"sep": ",","temperature": 0.2,"top_p": None,"num_beams": 1,'max_new_tokens': 512})()# 加载模型tokenizer, model, image_processor, context_len = load_model(model_path)# 运行模型评估try:response = eval_model(args)except Exception as e:response = f"Error: {str(e)}"# 将生成的数据保存到总文件中out_f.write(f"Image: {image_name}\n")out_f.write(f"Response: {response}\n")out_f.write("\n")# 这里我已经尝试清理显存了，下面还有很多不同的方法del modeltorch.cuda.empty_cache()

1. 及时释放显存

在循环中使用模型后，应该及时释放显存。
比如for循环中每次调用大模型一次结束时（我这里写的简单了，大家可以更深入的去研究一下）

# 清理显存
del model
torch.cuda.empty_cache()

torch.cuda.empty_cache() 可以释放未使用的显存，但不会影响已经分配的显存块。它可以减少显存碎片化，提高显存的利用效率。

2. 使用上下文管理器

使用上下文管理器可以确保在退出上下文时自动释放资源。对于模型加载和推理过程，可以使用上下文管理器来确保显存的及时释放：

from contextlib import contextmanager@contextmanager
def load_and_unload_model(model_path):tokenizer, model, image_processor, context_len = load_model(model_path)try:yield tokenizer, model, image_processor, context_lenfinally:del modeltorch.cuda.empty_cache()

然后在循环中使用这个上下文管理器：

for image in data:with load_and_unload_model(model_path) as (tokenizer, model, image_processor, context_len):# 运行模型评估try:response = eval_model(args)except Exception as e:response = f"Error: {str(e)}"# 将生成的数据保存到总文件中out_f.write(f"Image: {image_name}\n")out_f.write(f"Response: {response}\n")out_f.write("\n")

3. 分批处理数据

如果数据量很大，可以考虑将数据分批处理，以减少单次处理所需的显存。可以将数据分成若干小批次，每次只处理一个小批次的数据：

batch_size = 10
for i in range(0, len(data), batch_size):batch_data = data[i:i + batch_size]for image in batch_data:with load_and_unload_model(model_path) as (tokenizer, model, image_processor, context_len):# 运行模型评估try:response = eval_model(args)except Exception as e:response = f"Error: {str(e)}"# 将生成的数据保存到总文件中out_f.write(f"Image: {image_name}\n")out_f.write(f"Response: {response}\n")out_f.write("\n")

4. 使用更小的模型

如果显存仍然不足，可以考虑使用更小的模型或量化模型。例如，你的代码中提到的 load_4bit=True 和 load_8bit=True 可以显著减少模型的显存占用：

tokenizer, model, image_processor, context_len = load_pretrained_model(model_path=model_path,model_base=None,model_name=get_model_name_from_path(model_path),load_4bit=True,  # 或者 load_8bit=True
)

5. 动态调整显存分配

PyTorch 提供了一些选项来动态调整显存分配，例如 torch.cuda.memory_allocated() 和 torch.cuda.memory_reserved() 可以帮助你监控显存使用情况，做出相应的调整。

6. 使用混合精度训练

混合精度训练可以减少显存占用，并且通常可以提高计算速度。可以使用 torch.cuda.amp 来实现：

from torch.cuda.amp import autocast, GradScalerscaler = GradScaler()for image in data:with load_and_unload_model(model_path) as (tokenizer, model, image_processor, context_len):with autocast():# 运行模型评估try:response = eval_model(args)except Exception as e:response = f"Error: {str(e)}"# 将生成的数据保存到总文件中out_f.write(f"Image: {image_name}\n")out_f.write(f"Response: {response}\n")out_f.write("\n")

总结

显存管理在深度学习中非常重要，通过及时释放显存、使用上下文管理器、分批处理数据、使用更小的模型、动态调整显存分配以及混合精度训练，可以有效地管理显存，减少显存不足的问题。

希望这些方法和技巧对你有所帮助。