前言，在pytorch中，當(dāng)服務(wù)器上的gpu被占用時(shí)，很多時(shí)候我們想先用cpu調(diào)試下代碼，那么就需要進(jìn)行g(shù)pu和cpu的切換。

方法1：x.to(device)

把 device 作為一個(gè)可變參數(shù)，推薦使用argparse進(jìn)行加載：

使用gpu時(shí)：

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

使用cpu時(shí)：

device='cpu'
x.to(device) 

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

很多貼子中說(shuō)，使用x.cuda() 和x.to('cuda') 雖然是等效的，但是x.cuda() 的缺點(diǎn)是無(wú)法動(dòng)態(tài)切換cpu。然而，其實(shí)配合命令行參數(shù)CUDA_VISIBLE_DEVICES 是可以進(jìn)行切換的。

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

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

首先先看一下，正常運(yùn)行的情況：

• 執(zhí)行命令：python t.py
• 輸出結(jié)果：因?yàn)榉?wù)器上有兩個(gè)gpu，所以是我們想要的結(jié)果。

2
True

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

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

1
True

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

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

0
False

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

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

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