pytorch使用多GPU

fancyNSEU

于 2021-12-25 13:17:04 发布

阅读量1.1k

点赞数 1

文章标签： pytorch 深度学习机器学习

本文链接：https://blog.csdn.net/yule200910/article/details/122142207

版权

PyTorch GPU 设备管理 DataParallel 模型并行

关键词由CSDN通过智能技术生成

本文介绍了如何在PyTorch中检查和利用GPU资源。通过`torch.cuda.device_count()`检查GPU数量，使用`nvidia-smi`获取显卡信息。在多GPU环境下，通过`torch.nn.DataParallel`实现模型并行训练，并将数据和模型分配到特定GPU。同时讨论了如何选择部分GPU进行训练。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

# 查看几块GPU
import torch
print(torch.cuda.device_count())

# 查看显卡配置信息
nvidia-smi

# 单主机多块GPU使用
device = torch.device("cuda:0" if torch.cuda.is_avaliable() else :"cpu")
device0 = torch.device("cuda:0")
device1 = torch.device("cuda:1")

net = torch.nn.DataParallel(model)  # 可用的GPU都可以被使用
net.to(device)

# 如何利用部分GPU ？？？？？
# 假设有4块GPU 【0,1,2,3】
device_id = [0, 1, 2, 3]
input_data = input_data.to(device=device_id[0])
net = torch.nn.DataParallel(model)
net.to(device)

Pytorch使用GPU进行训练注意事项 | 文艺数学君