这篇文章主要介绍了pytorch中.to(device) 和.cuda()的区别有哪些,具有一定借鉴价值,感兴趣的朋友可以参考下,希望大家阅读完这篇文章之后大有收获,下面让小编带着大家一起了解一下。
原理
.to(device) 可以指定CPU 或者GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 单GPU或者CPUmodel.to(device)#如果是多GPUif torch.cuda.device_count() > 1: model = nn.DataParallel(model,device_ids=[0,1,2])model.to(device)
.cuda() 只能指定GPU
#指定某个GPUos.environ['CUDA_VISIBLE_DEVICE']='1'model.cuda()#如果是多GPUos.environment['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'device_ids = [0,1,2,3]net = torch.nn.Dataparallel(net, device_ids =device_ids)net = torch.nn.Dataparallel(net) # 默认使用所有的device_ids net = net.cuda()
class DataParallel(Module): def __init__(self, module, device_ids=None, output_device=None, dim=0): super(DataParallel, self).__init__() if not torch.cuda.is_available(): self.module = module self.device_ids = [] return if device_ids is None: device_ids = list(range(torch.cuda.device_count())) if output_device is None: output_device = device_ids[0]
补充:Pytorch使用To方法编写代码在不同设备(CUDA/CPU)上兼容(device-agnostic)
以前版本的PyTorch编写device-agnostic代码非常困难(即,在不修改代码的情况下在CUDA可以使用或者只能使用CPU的设备上运行)。
device-agnostic的概念
即设备无关,可以理解为无论什么设备都可以运行您编写的代码。(PS:个人理解,我没有在网上找到专业解释)
PyTorch 0.4.0使代码兼容
PyTorch 0.4.0通过两种方法使代码兼容变得非常容易:
张量的device属性为所有张量提供了torch.device设备。(注意:get_device仅适用于CUDA张量)
to方法Tensors和Modules可用于容易地将对象移动到不同的设备(代替以前的cpu()或cuda()方法)
我们推荐以下模式:
# 开始脚本,创建一个张量device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")...# 但是无论你获得一个新的Tensor或者Module# 如果他们已经在目标设备上则不会执行复制操作input = data.to(device)model = MyModule(...).to(device)
感谢你能够认真阅读完这篇文章,希望小编分享的“pytorch中.to(device) 和.cuda()的区别有哪些”这篇文章对大家有帮助,同时也希望大家多多支持编程网,关注编程网行业资讯频道,更多相关知识等着你来学习!