大家好，我在使用Hugging Face Transformers加载facebook/vit-msn-base-4模型做推理时，碰到了两个棘手的内存异常问题，折腾了好久都没解决，来求助各位大佬：

问题1：CPU推理时内存占用过高且持续增长 #

我硬编码指定模型在CPU上运行（device='cpu'），这个模型本身是只有12层的ViT-MSN-base，文件大小才160MB左右，而且batch size设的是1，但Python进程居然占了大概3GB内存。更头疼的是，我已经在代码里尝试用gc.collect()做内存清理，但内存占用还是一直在缓慢上升，根本降不下来。

问题2：CUDA推理时内存暴增并触发OOM #

切换到GPU（device='cuda'）运行后，内存直接飙到了10GB左右，最后直接报了CUDA内存不足的错误。我查了下资料，有人说是PyTorch的内存碎片问题导致的，但具体该怎么处理完全没头绪。而且报错里显示PyTorch已经分配了10.17GB，但我的GPU总容量才4GB，这和任务管理器里显示的实际内存占用也对不上，太奇怪了。

相关代码片段 #

模型基础配置

model_name  = 'facebook/vit-msn-base-4'
layers = 1

模型加载与主逻辑

cka_score_list = []

tokenizer = AutoImageProcessor.from_pretrained(
    model_name,
    use_fast=True,             # 使用快速tokenizer实现
    trust_remote_code=True,    # 信任远程代码（部分模型需要）
    add_bos_token=False,       # 不添加序列开始token
    add_eos_token=False,       # 不添加序列结束token
    padding_side="left"        # 在左侧补全序列
)

# 加载预训练模型
model = AutoModel.from_pretrained(
    model_name,
    trust_remote_code=True,    # 信任远程代码
    device_map="auto",         # 自动将层分配到可用设备
    torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32,  # 支持的话用bfloat16
)
print(dir(model))
model.eval()
model = model.to(device)
cur_score = eval_meth(classA, classB, num_images)

cka_score_list.append(cur_score)

核心推理与内存清理函数

def latent_embeddings_for_class(classA, num_images):
    cur_csv = os.path.join(val_csv , (classA+ '.csv'))
    df = pd.read_csv(cur_csv, header = None)
    embeddings = []
    for i in range(num_images):
        image_name = df.iloc[random.randint(1, 50), 0]
        image_path = os.path.join(val_images, image_name)
        image = Image.open(image_path)

        inputs = tokenizer(image, return_tensors="pt").to(device).to(torch.bfloat16)
        print(inputs['pixel_values'].shape)
        snapshot = tracemalloc.take_snapshot()

        # 打印内存占用Top3的行
        top_stats = snapshot.statistics('lineno')
        for stat in top_stats[:3]:
            print(stat)

        output = model(**inputs)
        embeddings.append(output.last_hidden_state[0,1])

        del inputs
        del image
        del output
        clear_memory()
        
        print(embeddings[0].shape)
    
    embeddings = torch.stack(embeddings)
    return embeddings

def eval_meth(classA, classB, num_images ):
    X= latent_embeddings_for_class( classA, num_images)
    Y= latent_embeddings_for_class( classB, num_images)
    
    cka_score = cka(X,Y)
    clear_memory()
    return cka_score

相关输出信息 #

CPU运行时的tracemalloc输出

torch.Size([1, 3, 224, 224])
<frozen importlib._bootstrap_external>:647: size=3355 KiB, count=19535, average=176 B
c:\Users\hp\miniconda3\envs\mka_research\lib\ast.py:50: size=661 KiB, count=11348, average=60 B
c:\Users\hp\miniconda3\envs\mka_research\lib\selectors.py:315: size=288 KiB, count=6, average=48.0 KiB
torch.Size([768])
torch.Size([1, 3, 224, 224])
<frozen importlib._bootstrap_external>:647: size=3355 KiB, count=19535, average=176 B
c:\Users\hp\miniconda3\envs\mka_research\lib\ast.py:50: size=661 KiB, count=11348, average=60 B
c:\Users\hp\miniconda3\envs\mka_research\lib\tracemalloc.py:505: size=537 KiB, count=9815, average=56 B
torch.Size([768])
torch.Size([1, 3, 224, 224])
<frozen importlib._bootstrap_external>:647: size=3355 KiB, count=19535, average=176 B
c:\Users\hp\miniconda3\envs\mka_research\lib\ast.py:50: size=661 KiB, count=11348, average=60 B
c:\Users\hp\miniconda3\envs\mka_research\lib\tracemalloc.py:505: size=533 KiB, count=9732, average=56 B
torch.Size([768])

GPU运行时的OOM报错

CUDA out of memory. Tried to allocate 452.00 MiB. GPU 0 has a total capacity of 4.00 GiB of which 0 bytes is free. Of the allocated memory 10.17 GiB is allocated by PyTorch, and 144.56 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

备注：内容来源于stack exchange，提问作者shivam