這篇文章主要介紹了pytorch如何使用指定GPU訓練，具有一定借鑒價值，感興趣的朋友可以參考下，希望大家閱讀完這篇文章之后大有收獲，下面讓小編帶著大家一起了解一下。
雖然pytorch提供了指定gpu的幾種方式，但是使用不當的話會遇到out of memory的問題，主要是因為pytorch會在第0塊gpu上初始化，并且會占用一定空間的顯存。這種情況下，經常會出現指定的gpu明明是空閑的，但是因為第0塊gpu被占滿而無法運行，一直報out of memory錯誤。
解決方案如下：
指定環境變量，屏蔽第0塊gpu
CUDA_VISIBLE_DEVICES = 1 main.py
這句話表示只有第1塊gpu可見，其他gpu不可用，此時要注意第1塊gpu已經變成第0塊，因此代碼里應該使用編號0來指定gpu。如果依然使用cuda:1會報invalid device ordinal。
當然也可以在代碼里使用os模塊達到同樣效果：
import os #多塊使用逗號隔開 os.environ['CUDA_VISIBLE_DEVICES'] = '1'
使用torch.cuda.device with torch.cuda.device(1): ...
注意這種情況下gpu編號也會改變
感謝你能夠認真閱讀完這篇文章，希望小編分享的“pytorch如何使用指定GPU訓練”這篇文章對大家有幫助，同時也希望大家多多支持億速云，關注億速云行業資訊頻道，更多相關知識等著你來學習!
免責聲明：本站發布的內容（圖片、視頻和文字）以原創、轉載和分享為主，文章觀點不代表本網站立場，如果涉及侵權請聯系站長郵箱：is@yisu.com進行舉報，并提供相關證據，一經查實，將立刻刪除涉嫌侵權內容。
