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)
#指定某個(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)行）。
即設(shè)備無(wú)關(guān)，可以理解為無(wú)論什么設(shè)備都可以運(yùn)行您編寫的代碼。(PS:個(gè)人理解，我沒(méi)有在網(wǎng)上找到專業(yè)解釋)
PyTorch 0.4.0通過(guò)兩種方法使代碼兼容變得非常容易：
張量的device屬性為所有張量提供了torch.device設(shè)備。（注意：get_device僅適用于CUDA張量）
to方法Tensors和Modules可用于容易地將對(duì)象移動(dòng)到不同的設(shè)備（代替以前的cpu()或cuda()方法）
我們推薦以下模式：
# 開(kāi)始腳本，創(chuàng)建一個(gè)張量
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
...
# 但是無(wú)論你獲得一個(gè)新的Tensor或者M(jìn)odule
# 如果他們已經(jīng)在目標(biāo)設(shè)備上則不會(huì)執(zhí)行復(fù)制操作
input = data.to(device)
model = MyModule(...).to(device)
以上為個(gè)人經(jīng)驗(yàn)，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。
標(biāo)簽：成都 江蘇 常州 六盤水 宿遷 駐馬店 山東 蘭州
巨人網(wǎng)絡(luò)通訊聲明：本文標(biāo)題《pytorch中.to(device) 和.cuda()的區(qū)別說(shuō)明》，本文關(guān)鍵詞 pytorch,中,.to,device,和,.cuda,；如發(fā)現(xiàn)本文內(nèi)容存在版權(quán)問(wèn)題，煩請(qǐng)?zhí)峁┫嚓P(guān)信息告之我們，我們將及時(shí)溝通與處理。本站內(nèi)容系統(tǒng)采集于網(wǎng)絡(luò)，涉及言論、版權(quán)與本站無(wú)關(guān)。
