在Docker for PyTorch模型中访问GPU的方式可以通过以下步骤来实现:
--gpus
参数来指定要将GPU暴露给容器。例如,--gpus all
将暴露宿主机上的所有GPU。import torch
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
# 在模型创建之前设置默认设备为GPU
torch.cuda.set_device(device)
model = MyModel().to(device)
data = data.to(device)
output = model(data)
这样,您的PyTorch模型将能够在Docker容器中访问并利用GPU进行高性能计算。
关于Docker和PyTorch的更多详细信息,请参考以下链接:
领取专属 10元无门槛券
手把手带您无忧上云