pytorch中.to(device) 和.cuda()的区别说明
原理
.to(device) 可以指定CPU 或者GPU
1 2 3 4 5 6 | device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu" ) # 单GPU或者CPU model.to(device) #如果是多GPU if torch.cuda.device_count() > 1 : model = nn.DataParallel(model,device_ids = [ 0 , 1 , 2 ]) model.to(device) |
.cuda() 只能指定GPU
1 2 3 4 5 6 7 8 9 | #指定某个GPU os.environ[ 'CUDA_VISIBLE_DEVICE' ] = '1' model.cuda() #如果是多GPU os.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() 采用.cuda的形式那么我们必须事先通过系统内部函数os.environment[ 'CUDA_VISIBLE_DEVICES' ] 获取到系统的GPU |
1 2 3 4 5 6 7 8 9 10 11 12 13 | 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的设备上运行)。
PyTorch 0.4.0使代码兼容
PyTorch 0.4.0通过两种方法使代码兼容变得非常容易:
张量的device属性为所有张量提供了torch.device设备。(注意:get_device仅适用于CUDA张量)
to方法Tensors和Modules可用于容易地将对象移动到不同的设备(代替以前的cpu()或cuda()方法)
我们推荐以下模式:
1 2 3 4 5 6 7 | # 开始脚本,创建一个张量 device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu" ) ... # 但是无论你获得一个新的Tensor或者Module # 如果他们已经在目标设备上则不会执行复制操作 input = data.to(device) model = MyModule(...).to(device) |