Pytorch如何指定device(cuda or cpu)
Pytorch指定device(cuda or cpu)
torch.device()
在PyTorch中,你可以使用torch.device
函數(shù)來設(shè)置設(shè)備(device)。
這個函數(shù)接受一個字符串作為參數(shù),該字符串可以是 "cpu" 或 "cuda",它指定了要在哪個設(shè)備上運行張量計算。
# 使用CPU device = torch.device("cpu") # 使用GPU(如果可用) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cuda:0") device = torch.device("cuda:1")
os.environ['CUDA_VISIBLE_DEVICES']
全局指定只能識別的cuda編號
os.environ['CUDA_VISIBLE_DEVICES']
是一個在運行 Python 程序時可以設(shè)置可見的 CUDA 設(shè)備的環(huán)境變量。
這個環(huán)境變量可以控制哪些 GPU 設(shè)備對 Python 程序是可見的。
例如:
如果你想讓你的程序只能看到編號為 0 的 GPU,你可以在運行程序前設(shè)置這個環(huán)境變量:
import os os.environ['CUDA_VISIBLE_DEVICES'] = '0'
或者你也可以設(shè)置多個 GPU,只需將它們的編號以逗號分隔:
import os os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
如果你想讓程序?qū)λ?GPU 都可見,只需將環(huán)境變量設(shè)置為空字符串:
os.environ['CUDA_VISIBLE_DEVICES'] = ''
說說torch.device的作用
torch.device代表將torch.Tensor分配到的設(shè)備的對象。
torch.device包含一個設(shè)備類型(‘cpu’或‘cuda’)和可選的設(shè)備序號。
如果設(shè)備序號不存在,則為當(dāng)前設(shè)備。
如:torch.Tensor用設(shè)備構(gòu)建‘cuda’的結(jié)果等同于‘cuda:X’,其中X是torch.cuda.current_device()的結(jié)果。
#開始腳本,創(chuàng)建一個張量 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") ... #但是無論你獲得一個新的Tensor或者M(jìn)odule #如果他們已經(jīng)在目標(biāo)設(shè)備上則不會執(zhí)行復(fù)制操作 input = data.to(device) model = MyModule(...).to(device)
總結(jié)
以上為個人經(jīng)驗,希望能給大家一個參考,也希望大家多多支持腳本之家。
相關(guān)文章
使用pipenv管理python虛擬環(huán)境的全過程
pipenv 是Kenneth Reitz大神的作品,能夠有效管理Python多個環(huán)境,各種包,接下來通過本文給大家分享使用pipenv管理python虛擬環(huán)境的全過程,感興趣的朋友一起看看吧2021-09-09Python使用pyinstaller打包spec文件的方法詳解
PyInstaller是一個用于將Python腳本打包成獨立的可執(zhí)行文件的工具,使用PyInstaller您可以將Python應(yīng)用程序轉(zhuǎn)換為可執(zhí)行文件,而無需用戶安裝Python解釋器或任何額外的庫,這篇文章主要給大家介紹了關(guān)于Python使用pyinstaller打包spec文件的相關(guān)資料,需要的朋友可以參考下2024-08-08python實現(xiàn)網(wǎng)絡(luò)五子棋
這篇文章主要為大家詳細(xì)介紹了python實現(xiàn)網(wǎng)絡(luò)五子棋,文中示例代碼介紹的非常詳細(xì),具有一定的參考價值,感興趣的小伙伴們可以參考一下2021-04-04python統(tǒng)計中文字符數(shù)量的兩種方法
今天小編就為大家分享一篇python統(tǒng)計中文字符數(shù)量的兩種方法,具有很好的參考價值,希望對大家有所幫助。一起跟隨小編過來看看吧2019-01-01