如何将PyTorch模型（Whisper Large V2）直接加载至GPU

你遇到的问题是PyTorch在反序列化模型时会先将整个张量存储加载到CPU内存，再转移到GPU，导致需要额外的大内存虚拟机。以下是两种可行的解决方案：

方案一：使用torch.load的mmap参数（OpenAI原生Whisper适用） #

PyTorch 1.12+支持mmap=True参数，通过内存映射文件的方式读取模型，避免将整个模型加载到CPU内存，直接将张量数据读取并转移到GPU。修改你的代码如下：

from whisper import Whisper, ModelDimensions
import torch

checkpoint_file = "large-v2.pt"
# 启用内存映射，指定直接加载到GPU
checkpoint = torch.load(checkpoint_file, map_location="cuda", mmap=True)

dims = ModelDimensions(**checkpoint["dims"])
model = Whisper(dims)
# 加载state_dict时，张量已经在GPU上，无需额外CPU内存
model.load_state_dict(checkpoint["model_state_dict"])

# 清理临时变量释放资源
del checkpoint

原理说明：mmap=True让PyTorch直接从磁盘映射模型文件，张量数据不会被完整拷贝到CPU内存，而是按需读取并直接分配到指定的GPU设备，大幅降低CPU内存占用。

方案二：改用Hugging Face Transformers库的Whisper实现 #

如果可以切换到Hugging Face的实现，其内置的内存优化机制支持直接将模型加载到GPU，无需占用大量CPU内存：

from transformers import WhisperForConditionalGeneration

# 加载模型时自动分配到GPU，无需先加载到CPU
model = WhisperForConditionalGeneration.from_pretrained(
    "openai/whisper-large-v2",  # 也可以指定本地模型路径
    device_map="auto"
)

原理说明：该方法通过accelerate库的内存优化策略，自动将模型各层分配到GPU，加载过程中不会在CPU内存中缓存整个模型，适合云端资源受限的场景。

注意事项 #

• 确保使用PyTorch 1.12及以上版本，mmap参数在旧版本中可能不稳定或不支持。
• 云端环境建议使用SSD存储模型文件，提升内存映射的读取速度。
• 使用OpenAI原生Whisper时，mmap是最轻量化的修改方式，无需更换依赖库。

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