pytorch 中的 model.cuda()

数据挖掘 火炬
2021-09-15 06:37:55

如果我model.cuda()在 pytorch 中调用 model 的子类nn.Module,并说如果我有四个 GPU,它将如何利用这四个 GPU,我如何知道哪些 GPU 正在使用?

1个回答

model.cuda()默认情况下会将您的模型发送到“当前设备”,可以使用torch.cuda.set_device(device).

将模型发送到特定设备的另一种方法是model.to(torch.device('cuda:0')).

当然,这取决于环境变量中指定的设备可见性CUDA_VISIBLE_DEVICES

您可以使用nvidia-smi. 此外,nvtop这非常好。

PyTorch 在多个 GPU 中训练模型的标准方法是使用nn.DataParallel将模型复制到 GPU 并在训练期间在它们之间拆分批次并组合各个输出。