前言，在pytorch中，當服務器上的gpu被占用時，很多時候我們想先用cpu調試下代碼，那么就需要進行gpu和cpu的切換。

方法1：x.to(device)

把 device 作為一個可變參數，推薦使用argparse進行加載：

使用gpu時：

device='cuda'
x.to(device) # x是一個tensor，傳到cuda上去

使用cpu時：

device='cpu'
x.to(device) 

方法2：使用x.cuda()+CUDA_VISIBLE_DEVICES

很多貼子中說，使用x.cuda() 和x.to('cuda') 雖然是等效的，但是x.cuda() 的缺點是無法動態切換cpu。然而，其實配合命令行參數CUDA_VISIBLE_DEVICES 是可以進行切換的。

在服務器上創建一個python腳本 t.py：

import torch
print(torch.cuda.device_count()) # 可用gpu數量
print(torch.cuda.is_available()) # 是否可用gpu

首先先看一下，正常運行的情況：

• 執行命令：python t.py
• 輸出結果：因為服務器上有兩個gpu，所以是我們想要的結果。

2
True

如果想要只使用某一塊gpu，只需要在執行前加一個參數：

• CUDA_VISIBLE_DEVICES=0 python t.py，例如，我們要使用gpu 0
• 接下來看看輸出什么：是的！程序中確實只可見了一塊gpu~
  

1
True

下面，如果我們想使用cpu呢？

• CUDA_VISIBLE_DEVICES="" python t.py
• 輸出結果：可以看到，雖然服務器上有2塊cpu，通過我們設置執行參數，程序中也成功看不到了！
  

0
False

因此，回歸正題，當我們使用x.cuda()進行分配gpu時，只需要使用torch.cuda.is_available()加一個判斷即可，當想使用cpu的時候在執行程序的命令行參數進行控制：

if torch.cuda.is_available():
  x= x.cuda()

到此這篇關于Pytorch如何切換 cpu和gpu的使用詳解的文章就介紹到這了,更多相關Pytorch切換cpu和gpu內容請搜索腳本之家以前的文章或繼續瀏覽下面的相關文章希望大家以后多多支持腳本之家！