怎么在Pytorch中切換 cpu和gpu

本篇文章給大家分享的是有關怎么在Pytorch中切換 cpu和gpu，小編覺得挺實用的，因此分享給大家學習，希望大家閱讀完這篇文章后可以有所收獲，話不多說，跟著小編一起來看看吧。

方法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，小編相信有部分知識點可能是我們日常工作會見到或用到的。希望你能通過這篇文章學到更多知識。更多詳情敬請關注億速云行業資訊頻道。