CUDA语义


torch.cuda会跟踪当前选择的GPU,并且将分配您的所有CUDA张量。可以使用torch.cuda.device上下文管理器更改所选设备 。

但是,一旦张量被分配,您可以对其进行操作,而不考虑您选择的设备,结果将始终放在与张量相同的设备上。

除了copy_()以外,默认是不支持跨GPU操作。除非启用对等存储器访问,否则在不同设备上分布的张量上尝试启动的操作都会引发错误。

下面你可以找到一个展示如下的小例子:

x = torch.cuda.FloatTensor(1)
# x.get_device() == 0
y = torch.FloatTensor(1).cuda()
# y.get_device() == 0

with torch.cuda.device(1):
    # allocates a tensor on GPU 1
    a = torch.cuda.FloatTensor(1)

    # transfers a tensor from CPU to GPU 1
    b = torch.FloatTensor(1).cuda()
    # a.get_device() == b.get_device() == 1

    c = a + b
    # c.get_device() == 1

    z = x + y
    # z.get_device() == 0

    # even within a context, you can give a GPU id to the .cuda call
    d = torch.randn(2).cuda(2)
    # d.get_device() == 2

最佳实践

使用固定的内存缓冲区

当它们来自固定(页锁)存储器时,主机到GPU副本的速度要快得多。CPU张量和存储器暴露一种pin_memory() 方法,返回对象的副本,数据放在固定区域中。

另外,一旦固定了张量或存储空间,就可以使用异步GPU副本。只需传递一个附加async=True参数给一个cuda() 调用。这可以用于将数据传输与计算重叠。

您可以DataLoader通过传递pin_memory=True给其构造函数将返回批处理置于固定内存中。

使用 nn.DataParallel 替代 multiprocessing

大多数涉及批量输入和多个GPU的情况默认使用DataParallel来使用多个GPU。即使使用GIL,单个python进程也可能使多个GPU饱和。

从0.1.9版本开始,大量的GPU(8+)可能未被充分利用。然而,这是一个已知的问题,也正在积极开发。和往常一样,测试你的用例吧。

使用CUDA模型调用multiprocessing需要特别注意;除非需要谨慎地满足数据处理需求,否则您的程序很可能会出现错误或未定义的行为。