Pytorch调用GPU的方法

Pytorch调用GPU有两种方法，一种是torch.cuda()，一种是torch.to()。

torch.cuda()

通常会在配置文件中写入调用的GPU，默认不填参数为0
gpu = [0, 1, 2]
# 默认调用调用0号GPU
network = network.cuda()
# 指定调用0号GPU
network = network.cuda(gpu[0])
# 指定调用1号GPU
network = network.cuda(gpu[1])
1
2
3
4
5
6
7
8

torch.to()

GPU = ['cuda', 'cuda:0', 'cuda:1', 'cuda:2']
# 这个写法是最常见的写法，默认调用0号GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 指定调用0号GPU
device = torch.device(GPU[1] if torch.cuda.is_available() else "cpu")
# 指定调用1号GPU
device = torch.device(GPU[2] if torch.cuda.is_available() else "cpu")

network = network.to(device)

1
2
3
4
5
6
7
8
9
10