我来帮你搞定这个Flask API的GPU显存管理问题——你遇到的显存泄漏其实是PyTorch内存回收机制和Flask请求上下文结合的典型问题，下面给你几个实操性强的解决方案：

你遇到的GPU内存持续占用，本质上是这几个原因叠加导致的：

• 推理生成的张量（包括中间结果、输出）仍驻留在GPU，没有被主动清理
• Python的垃圾回收（GC）不会立即回收GPU张量的引用
• Flask的多线程/请求上下文可能会持有模型或张量的引用，导致无法被回收
• 没有禁用梯度计算，推理时产生了不必要的梯度张量占用显存

1. 推理后强制清理GPU资源（最关键的一步） #

每次推理完成后，不仅要把模型移回CPU，还要主动清理所有相关的GPU张量、触发GC并清空CUDA缓存。这里给你一段可复用的代码：

import gc
import torch
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer

def load_model(model_name):
    # 自定义模型加载逻辑，示例为分类模型
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    return model, tokenizer

def run_inference(model_name, inputs):
    # 加载模型到CPU（先放CPU，避免加载时直接占GPU）
    model, tokenizer = load_model(model_name)
    # 移到GPU执行推理
    model = model.to("cuda")
    
    # 禁用梯度计算，避免生成冗余梯度张量
    with torch.no_grad():
        pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, device=0)
        outputs = pipe(inputs)
    
    # 1. 把模型移回CPU
    model = model.to("cpu")
    # 2. 递归清理所有输出中的GPU张量
    def move_tensor_to_cpu(obj):
        if isinstance(obj, torch.Tensor):
            return obj.cpu()
        elif isinstance(obj, dict):
            return {k: move_tensor_to_cpu(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [move_tensor_to_cpu(item) for item in obj]
        else:
            return obj
    outputs = move_tensor_to_cpu(outputs)
    # 3. 删除所有模型相关引用
    del model, tokenizer, pipe
    # 4. 强制触发Python垃圾回收
    gc.collect()
    # 5. 清空CUDA缓存（释放CUDA层面未被回收的显存）
    torch.cuda.empty_cache()
    
    return outputs

重点说明：torch.cuda.empty_cache()是释放CUDA驱动层面保留的显存，而gc.collect()是清理Python层面的无效引用，两者结合才能彻底释放GPU内存。

2. 用全局状态管理模型加载，避免重复加载 #

维护一个全局变量跟踪当前在GPU的模型，切换模型时先清理旧模型，再加载新模型，减少重复加载的开销：

# 全局变量：跟踪当前加载的模型和名称
current_model = None
current_tokenizer = None
current_model_name = None

def get_model(model_name):
    global current_model, current_tokenizer, current_model_name
    # 如果是同一个模型，直接移到GPU复用
    if current_model_name == model_name:
        return current_model.to("cuda"), current_tokenizer
    # 切换模型前，先清理旧模型的资源
    if current_model is not None:
        current_model.to("cpu")
        del current_model, current_tokenizer
        gc.collect()
        torch.cuda.empty_cache()
    # 加载新模型到CPU，再移到GPU
    current_model, current_tokenizer = load_model(model_name)
    current_model_name = model_name
    return current_model.to("cuda"), current_tokenizer

# 推理函数修改为复用全局模型
def run_inference_optimized(model_name, inputs):
    model, tokenizer = get_model(model_name)
    with torch.no_grad():
        pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, device=0)
        outputs = pipe(inputs)
    # 推理完成后把模型移回CPU（保留引用，下次复用）
    model.to("cpu")
    # 清理输出张量
    outputs = move_tensor_to_cpu(outputs)
    del pipe
    gc.collect()
    torch.cuda.empty_cache()
    return outputs

3. 限制Flask的并发请求数 #

Flask默认是多线程模式，如果同时有多个请求进来，可能会同时尝试加载多个模型到GPU，直接导致显存溢出。所以要强制单线程运行：

• 用Flask自带命令：flask run --threaded=False
• 用Gunicorn部署时：gunicorn --workers=1 --threads=1 app:app

因为你的需求是串行加载模型，单线程可以避免并发带来的显存冲突。

4. 检查并清理隐式张量引用 #

有些HuggingFace Pipeline内部会缓存中间结果（比如last_hidden_state），这些张量可能留在GPU上。如果发现显存还是没释放，可以尝试直接重置Pipeline的模型引用：

del pipe.model

可以在每次操作后打印GPU显存使用情况，确认是否真的释放：

print(f"已分配显存: {torch.cuda.memory_allocated()/1024**2:.2f} MB")
print(f"已保留显存: {torch.cuda.memory_reserved()/1024**2:.2f} MB")

如果清理后已分配显存回到接近初始值，说明释放成功。

内容的提问来源于stack exchange，提问作者OG0