Pytorch如何切換 cpu和gpu的使用詳解

Posted on by

前言，在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內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet！

推薦閱讀：