[英]A better way to make pytorch code agnostic to running on a CPU or GPU?
遷移指南建議以下內容使代碼CPU / GPU不可知:
> # at beginning of the script
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
...
# then whenever you get a new Tensor or Module
# this won't copy if they are already on the desired device
input = data.to(device)
model = MyModule(...).to(device)
我做了這個並在僅CPU設備上運行我的代碼,但我的模型在輸入一個輸入數組時崩潰了,因為它說它期望CPU張量不是GPU。 不知何故,我的模型自動將CPU輸入數組轉換為GPU數組。 最后,我在我的代碼中追溯到這個命令:
model = torch.nn.DataParallel(model).to(device)
即使我將模型轉換為'cpu',nn.DataParallel也會覆蓋它。 我想出的最佳解決方案是有條件的:
if device.type=='cpu':
model = model.to(device)
else:
model = torch.nn.DataParallel(model).to(device)
這似乎並不優雅。 有沒有更好的辦法?
怎么樣
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model = model.to(device)
?
如果只有一個GPU,則不需要DataParallel
。
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.