BLIP2模型加载在不同设备上

本文主要是介绍BLIP2模型加载在不同设备上，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

背景

现在大语言模型越来越大，占用的内存越来越多，这导致内存较小的设备无法体验大模型的效果。transformer提供了将一个大模型分别加载在gpu和cpu上的方法。

加载方法

1. 以多模态模型BLIP2为例，将其语言模型放在gpu上，其余部分放在cpu上。配置加载预加载模型的device_map.device_map可以设置为auto，则根据设备的显存情况，自动加载在gpu或者cpu上。
2. 使用BLIP2模型地址

device_map = {'language_model':0,\'language_projection':'cpu', \'qformer':'cpu', \'query_tokens':'cpu', \'vision_model':'cpu'}
multimodal = Blip2ForConditionalGeneration.from_pretrained("huggingface_opt",torch_dtype=torch.float16,device_map=device_map)

如上所示，将language_model放在0号显卡上，其余模型放在cpu上。模型的名称在模型的index文件中

测试情况

加载位置	显存占用	耗时	备注
GPU	13G	0.5s
GPU +CPU	10G	2s	语言放在0号显卡上，其余模型放在cpu上

这篇关于BLIP2模型加载在不同设备上的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---