pytorch 之cuda语义
torch.cuda 会记录当前选择的GPU,并且分配的所有CUDA张量将在上面创建。可以使用torch.cuda.device 上下文管理器更改所选设备。
但是,一旦张量被分配,您可以直接对其进行操作,而不考虑所选择的设备,结果将始终放在与张量相同的设备上
默认情况下,不支持跨GPU操作,唯一的例外是 copy_() 。 除非启用对等存储器访问,否则对分布不同设备上的张量任何启动操作的尝试都将会引发错误
import torch x = torch.cuda.FloatTensor(1) y = torch.FloatTensor(1).cuda() # print(x) #tensor([0.], device='cuda:0') # print(x.get_device()) #0 # print(y) #tensor([2.3332e-40], device='cuda:0') # print(y.get_device()) #0 with torch.cuda.device(1): a = torch.cuda.FloatTensor(1) #分配一个张量到 GPU 1 b = torch.FloatTensor(1).cuda() 把cpu上的张量转化到gpu上 c = a+b #a.get_device==1 and b.get_device==1 >>c.get_device==1 print(c.get_device()) #1 ,因为a.gpu=1,b.gpu=1,所以a+b的结果存放在gpu1上 z = x+y print(z.get_device()) #0,同上 d = torch.randn(2).cuda(1) #even within a context, you can give a GPU id to the .cuda call #对于一个变量,可以指定gpu print(d.get_device())#1
使用固定的内存缓冲区当副本来自固定(页锁)内存时,主机到GPU的复制速度要快很多。CPU张量和存储开放了一个 pin_memory() 方法,它返回该对象的副本,而它的数据放在固定区域中。另外,一旦固定了张量存储,就可以使用异步的GPU副本。只需传递一个额外的async=True 参数到 cuda() 的调用。这可以用于将数据传输与计算重叠。
通过将 pin_memory=True 传递给其构造函数,可以使 DataLoader 将batch返回到固定内存中 。
单GPU的使用
import os os.environ["CUDA_VISIBLE_DEVICES"] = "0"
多GPU的使用
device_ids = [0,1,2,3] model = model.cuda(device_ids[0]) model = nn.DataParallel(model, device_ids=device_ids)