pytorch使用GPU

1   单GPU

pytorch使用GPU，需要把数据和模型从内存转移到GPU

对于张量，一律使用方法 .to(device)  或者 .cuda()即可。  pytorch.Tensor(shapes).to(device)或者 torch.Tensor(shapes).cuda()

对于模型，也是使用同样的方法，使用.to(device) 或者.cuda()来将网络放到GPU显存。model.to(device)  或者  model.cuda()

 

其中devices=troch.device("cuda:1" if touch.cuda is_available() else  "cpu")

 

2  多GPU

单机多GPUs主要采用DataParallel函数，而不是DistributedParallel，后者一般用于多主机多GPUs，当然也可用于单机多GPU

使用多卡训练的方法很多，当然前提是我们的设备存在两个及以上的GPU

使用直接用model传入torch.nn.DataParallel函数即可。

device_ids=[1,2,3,4]

input_data=input_data.to(devices=device_ids[0])

对于模型

net=torch.nn.DataParallel(model)

net.to(device)