App下載

pytorch中to(device) 和cuda()有什么區(qū)別?如何使用?

猿友 2021-07-14 14:59:12 瀏覽數(shù) (6572)
反饋

很多小伙伴在使用 pytorch 指定驅(qū)動(dòng)的時(shí)候會(huì)遇到兩個(gè)功能類似的方法,就是?to(device)?和?cuda()?,這兩種方法的功能都是指定 CPU 或 GPU 進(jìn)行數(shù)據(jù)處理的,那么 pytorch 的 ?to(device)? 和 ?cuda() ?有什么區(qū)別呢?pytorch 該怎么用 ?to(device)? 和 ?cuda()? 呢?接下來的這篇文章告訴你。

原理

.to(device) 可以指定CPU 或者GPU

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

#指定某個(gè)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) # 默認(rèn)使用所有的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]

補(bǔ)充:Pytorch使用To方法編寫代碼在不同設(shè)備(CUDA/CPU)上兼容(device-agnostic)

以前版本的PyTorch編寫device-agnostic代碼非常困難(即,在不修改代碼的情況下在CUDA可以使用或者只能使用CPU的設(shè)備上運(yùn)行)。

device-agnostic的概念

即設(shè)備無關(guān),可以理解為無論什么設(shè)備都可以運(yùn)行您編寫的代碼。(PS:個(gè)人理解,我沒有在網(wǎng)上找到專業(yè)解釋)

PyTorch 0.4.0使代碼兼容

PyTorch 0.4.0通過兩種方法使代碼兼容變得非常容易:

張量的device屬性為所有張量提供了torch.device設(shè)備。(注意:get_device僅適用于CUDA張量)

to方法Tensors和Modules可用于容易地將對象移動(dòng)到不同的設(shè)備(代替以前的cpu()或cuda()方法)

我們推薦以下模式:

# 開始腳本,創(chuàng)建一個(gè)張量
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
...
# 但是無論你獲得一個(gè)新的Tensor或者M(jìn)odule
# 如果他們已經(jīng)在目標(biāo)設(shè)備上則不會(huì)執(zhí)行復(fù)制操作
input = data.to(device)
model = MyModule(...).to(device)

小結(jié)

以上就是pytorch的to(device)和cuda()有什么區(qū)別的全部介紹,希望能給大家一個(gè)參考,也希望大家多多支持W3Cschool。


0 人點(diǎn)贊