如何使用特定显卡跑pyTorch

版权声明:本文为博主原创文章,未经博主允许不得转载。 https://blog.csdn.net/u010698086/article/details/80346177

如何使用特定显卡跑pyTorch

方法一:改变系统变量使得仅目标显卡可见

export CUDA_VISIBLE_DEVICES=0 #这里是要使用的GPU编号,正常的话是从0开始

然后再运行pyTorch即可

方法二:使用torch.cuda接口

#在生成网络对象之前:
torch.cuda.set_device(0)

方法三:使用多pytorch的并行GPU接口

net = torch.nn.DataParallel(model, device_ids=[0])

方法四:初始化模型时

net = Net.cuda(0)
阅读更多

没有更多推荐了,返回首页