Pytorch如何指定device(cuda or cpu)
Pytorch指定device(cuda or cpu)
torch.device()
在PyTorch中，你可以使用torch.device函數(shù)來(lái)設(shè)置設(shè)備（device）。
這個(gè)函數(shù)接受一個(gè)字符串作為參數(shù)，該字符串可以是 "cpu" 或 "cuda"，它指定了要在哪個(gè)設(shè)備上運(yùn)行張量計(jì)算。
# 使用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']
全局指定只能識(shí)別的cuda編號(hào)
os.environ['CUDA_VISIBLE_DEVICES']
是一個(gè)在運(yùn)行 Python 程序時(shí)可以設(shè)置可見(jiàn)的 CUDA 設(shè)備的環(huán)境變量。
這個(gè)環(huán)境變量可以控制哪些 GPU 設(shè)備對(duì) Python 程序是可見(jiàn)的。
例如：
如果你想讓你的程序只能看到編號(hào)為 0 的 GPU，你可以在運(yùn)行程序前設(shè)置這個(gè)環(huán)境變量：
import os os.environ['CUDA_VISIBLE_DEVICES'] = '0'
或者你也可以設(shè)置多個(gè) GPU，只需將它們的編號(hào)以逗號(hào)分隔：
import os os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
如果你想讓程序?qū)λ?GPU 都可見(jiàn)，只需將環(huán)境變量設(shè)置為空字符串：
os.environ['CUDA_VISIBLE_DEVICES'] = ''
說(shuō)說(shuō)torch.device的作用
torch.device代表將torch.Tensor分配到的設(shè)備的對(duì)象。
torch.device包含一個(gè)設(shè)備類型（‘cpu’或‘cuda’）和可選的設(shè)備序號(hào)。
如果設(shè)備序號(hào)不存在，則為當(dāng)前設(shè)備。
如：torch.Tensor用設(shè)備構(gòu)建‘cuda’的結(jié)果等同于‘cuda：X’，其中X是torch.cuda.current_device()的結(jié)果。
#開(kāi)始腳本，創(chuàng)建一個(gè)張量
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
...
#但是無(wú)論你獲得一個(gè)新的Tensor或者M(jìn)odule
#如果他們已經(jīng)在目標(biāo)設(shè)備上則不會(huì)執(zhí)行復(fù)制操作
input = data.to(device)
model = MyModule(...).to(device)總結(jié)
以上為個(gè)人經(jīng)驗(yàn)，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。
相關(guān)文章
這篇文章主要介紹了python3 flask 文件占用未釋放問(wèn)題,具有很好的參考價(jià)值,希望對(duì)大家有所幫助,如有錯(cuò)誤或未考慮完全的地方,望不吝賜教2023-11-11
這篇文章主要為大家詳細(xì)介紹了Python中的隱含特性，文中示例代碼介紹的非常詳細(xì)，具有一定的參考價(jià)值，感興趣的小伙伴們可以參考一下，希望能夠給你帶來(lái)幫助2022-01-01
這篇文章主要介紹了Python多線程下的變量問(wèn)題,由于GIL的存在,Python的多線程編程問(wèn)題一直是開(kāi)發(fā)者中的熱點(diǎn)話題,需要的朋友可以參考下2015-04-04
使用pipenv管理python虛擬環(huán)境的全過(guò)程
pipenv 是Kenneth Reitz大神的作品，能夠有效管理Python多個(gè)環(huán)境，各種包，接下來(lái)通過(guò)本文給大家分享使用pipenv管理python虛擬環(huán)境的全過(guò)程，感興趣的朋友一起看看吧2021-09-09
Python使用pyinstaller打包spec文件的方法詳解
PyInstaller是一個(gè)用于將Python腳本打包成獨(dú)立的可執(zhí)行文件的工具,使用PyInstaller您可以將Python應(yīng)用程序轉(zhuǎn)換為可執(zhí)行文件,而無(wú)需用戶安裝Python解釋器或任何額外的庫(kù),這篇文章主要給大家介紹了關(guān)于Python使用pyinstaller打包spec文件的相關(guān)資料,需要的朋友可以參考下2024-08-08
在Python中,通常我們不會(huì)直接判斷for循環(huán)是否正在執(zhí)行最后一次迭代,因?yàn)镻ython的for循環(huán)是基于可迭代對(duì)象的,它不知道也不關(guān)心迭代的內(nèi)部狀態(tài)（比如當(dāng)前是第幾次迭代）,但是,我們可以使用一些技巧來(lái)間接地實(shí)現(xiàn)這個(gè)需求,需要的朋友可以參考下2025-01-01
python實(shí)現(xiàn)網(wǎng)絡(luò)五子棋
這篇文章主要為大家詳細(xì)介紹了python實(shí)現(xiàn)網(wǎng)絡(luò)五子棋，文中示例代碼介紹的非常詳細(xì)，具有一定的參考價(jià)值，感興趣的小伙伴們可以參考一下2021-04-04
Python提供了ord和chr兩個(gè)內(nèi)置的函數(shù)，用于字符與ASCII碼之間的轉(zhuǎn)換。2008-09-09
python 移動(dòng)圖片到另外一個(gè)文件夾的實(shí)例
今天小編就為大家分享一篇python 移動(dòng)圖片到另外一個(gè)文件夾的實(shí)例，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。一起跟隨小編過(guò)來(lái)看看吧2019-01-01
python統(tǒng)計(jì)中文字符數(shù)量的兩種方法
今天小編就為大家分享一篇python統(tǒng)計(jì)中文字符數(shù)量的兩種方法，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。一起跟隨小編過(guò)來(lái)看看吧2019-01-01