pytorch、tensorflow-多gpu顯存使用問題（解決gpu 0 顯存佔滿時 使用其他gpu 報錯 out of memory問題）

問題情景：

服務器上有兩個gpu，編號爲0的gpu在跑任務，顯存完全佔滿。gpu 1 空閒。在想用gpu 1 跑任務時，報錯 CUDA：out of memory. 此時已經設置了 device = torch.device(“cuda:1”） ，嘗試了設置device = torch.device(“cpu”），仍然會報錯。

原因：

gpu 0 爲主gpu，運行pytorch程序時 仍會嘗試向其上寫入內容，導致報錯。

解決方案：

修改環境變量，讓python程序忽視 gpu 0
比如 我要只使用gpu 1：

CUDA_VISIBLE_DEVICES=1 python my_script.py

如需在python程序內設置，使用如下語句：

import os
os.environ['CUDA_VISIBLE_DEVICES']='1'

參考資料：

https://github.com/pytorch/pytorch/issues/3477