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

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

方法1：x.to(device)

把 device 作為一個可變參數(shù)，推薦使用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() 的缺點是無法動態(tài)切換cpu。然而，其實配合命令行參數(shù)CUDA_VISIBLE_DEVICES 是可以進行切換的。

在服務器上創(chuàng)建一個python腳本 t.py：

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

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

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

2
True

如果想要只使用某一塊gpu，只需要在執(zhí)行前加一個參數(shù)：

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

1
True

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

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

0
False

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

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

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