大家好,我是爱编程的喵喵。双985硕士毕业,现担任全栈工程师一职,热衷于将数据思维应用到工作与生活中。从事机器学习以及相关的前后端开发工作。曾在阿里云、科大讯飞、CCF等比赛获得多次Top名次。现为CSDN博客专家、人工智能领域优质创作者。
本文主要介绍了transformers加载模型时自动选择空闲的多个GPU的实战代码,希望能对使用huggingface transformers的同学们有所帮助。
1. 背景描述
在上篇博客transformers加载模型时自动选择空闲的单个GPU的实战代码中介绍了单个GPU的实战代码,但是部分模型参数量比较大,需要至少两块GPU。所以就需要自动选择多块空闲的GPU。
在经过了亲身的实践后,终于找到了可复现的实战代码,最终将详细的代码总结如下。希望对拥有多张GPU并进行自动化分配的同学们有所帮助。