如何在两个GPU上运行两个不同的神经网络?

时间:2020-03-06 22:20:59

标签: python tensorflow gpu

我有两个GPU,但是现在仅使用一个GPU进行模型训练,那么如何在另一个GPU上运行另一个神经网络?

1 个答案:

答案 0 :(得分:0)

如果您使用的是pytorch,并且已安装CUDA,则可以使用其.to()函数

pytorch文档上的一个简单示例显示:

>>> tensor = torch.randn(2, 2)  # Initially dtype=float32, device=cpu
>>> tensor.to(torch.float64)
tensor([[-0.5044,  0.0005],
        [ 0.3310, -0.0584]], dtype=torch.float64)

>>> cuda0 = torch.device('cuda:0')
>>> tensor.to(cuda0)
tensor([[-0.5044,  0.0005],
        [ 0.3310, -0.0584]], device='cuda:0')

>>> tensor.to(cuda0, dtype=torch.float64)
tensor([[-0.5044,  0.0005],
        [ 0.3310, -0.0584]], dtype=torch.float64, device='cuda:0')

>>> other = torch.randn((), dtype=torch.float64, device=cuda0)
>>> tensor.to(other, non_blocking=True)
tensor([[-0.5044,  0.0005],
        [ 0.3310, -0.0584]], dtype=torch.float64, device='cuda:0')

您可以在这里参考这些链接:

https://pytorch.org/docs/stable/tensors.html#torch.Tensor.to https://pytorch.org/docs/stable/cuda.html#torch.cuda.device