pytorch中.to(device) 和.cuda()的區(qū)別說明
發(fā)布日期：2022-04-12 15:55 | 文章來源：站長之家
原理
.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
#指定某個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()
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:個人理解，我沒有在網(wǎng)上找到專業(yè)解釋)
PyTorch 0.4.0使代碼兼容
PyTorch 0.4.0通過兩種方法使代碼兼容變得非常容易：
張量的device屬性為所有張量提供了torch.device設備。（注意：get_device僅適用于CUDA張量）
to方法Tensors和Modules可用于容易地將對象移動到不同的設備（代替以前的cpu()或cuda()方法）
我們推薦以下模式：
# 開始腳本，創(chuàng)建一個張量
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
...
# 但是無論你獲得一個新的Tensor或者Module
# 如果他們已經(jīng)在目標設備上則不會執(zhí)行復制操作
input = data.to(device)
model = MyModule(...).to(device)
以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持本站。
版權聲明：本站文章來源標注為YINGSOO的內(nèi)容版權均為本站所有，歡迎引用、轉載，請保持原文完整并注明來源及原文鏈接。禁止復制或仿造本網(wǎng)站，禁止在非www.sddonglingsh.com所屬的服務器上建立鏡像，否則將依法追究法律責任。本站部分內(nèi)容來源于網(wǎng)友推薦、互聯(lián)網(wǎng)收集整理而來，僅供學習參考，不代表本站立場，如有內(nèi)容涉嫌侵權，請聯(lián)系alex-e#qq.com處理。
相關文章
關注官方微信