# 查看几块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
最新推荐文章于 2024-02-02 00:24:42 发布