# torch.cuda
# torch.cuda
該包增加了對CUDA張量類型的支持,實現了與CPU張量相同的功能,但使用GPU進行計算。
它是懶惰的初始化,所以你可以隨時導入它,并使用`is_available()`來確定系統是否支持CUDA。
[CUDA語義](cuda.html)中有關于使用CUDA的更多細節。
```
torch.cuda.current_blas_handle()
```
返回cublasHandle\_t指針,指向當前cuBLAS句柄
```
torch.cuda.current_device()
```
返回當前所選設備的索引。
```
torch.cuda.current_stream()
```
返回一個當前所選的`Stream`
```
class torch.cuda.device(idx)
```
上下文管理器,可以更改所選設備。
**參數:**
- **idx** (*int*) – 設備索引選擇。如果這個參數是負的,則是無效操作。
```
torch.cuda.device_count()
```
返回可得到的GPU數量。
```
class torch.cuda.device_of(obj)
```
將當前設備更改為給定對象的上下文管理器。
可以使用張量和存儲作為參數。如果給定的對象不是在GPU上分配的,這是一個無效操作。
**參數:**
- **obj** (*Tensor* or *Storage*) – 在選定設備上分配的對象。
```
torch.cuda.is_available()
```
返回一個bool值,指示CUDA當前是否可用。
```
torch.cuda.set_device(device)
```
設置當前設備。
不鼓勵使用此函數來設置。在大多數情況下,最好使用`CUDA_VISIBLE_DEVICES`環境變量。
**參數:**
- **device** (*int*) – 所選設備。如果此參數為負,則此函數是無效操作。
```
torch.cuda.stream(stream)
```
選擇給定流的上下文管理器。
在其上下文中排隊的所有CUDA核心將在所選流上入隊。
**參數:**
- **stream** (*Stream*) – 所選流。如果是`None`,則這個管理器是無效的。
```
torch.cuda.synchronize()
```
等待當前設備上所有流中的所有核心完成。
## 交流集
```
torch.cuda.comm.broadcast(tensor, devices)
```
向一些GPU廣播張量。
**參數:**
- **tensor** (*Tensor*) – 將要廣播的張量
- **devices** (*Iterable*) – 一個可以廣播的設備的迭代。注意,它的形式應該像(src,dst1,dst2,...),其第一個元素是廣播來源的設備。
**返回:** 一個包含張量副本的元組,放置在與設備的索引相對應的設備上。
```
torch.cuda.comm.reduce_add(inputs, destination=None)
```
將來自多個GPU的張量相加。
所有輸入應具有匹配的形狀。
**參數:**
- **inputs** (*Iterable\[Tensor\]*) – 要相加張量的迭代
- **destination** (*int*, optional) – 將放置輸出的設備(默認值:當前設備)。
**返回:** 一個包含放置在`destination`設備上的所有輸入的元素總和的張量。
```
torch.cuda.comm.scatter(tensor, devices, chunk_sizes=None, dim=0, streams=None)
```
打散橫跨多個GPU的張量。
**參數:**
- **tensor** (*Tensor*) – 要分散的張量
- **devices** (*Iterable\[int\]*) – int的迭代,指定哪些設備應該分散張量。
- **chunk\_sizes** (*Iterable\[int\]*, optional) – 要放置在每個設備上的塊大小。它應該匹配`devices`的長度并且總和為`tensor.size(dim)`。 如果沒有指定,張量將被分成相等的塊。
- **dim** (*int*, optional) – 沿著這個維度來chunk張量
**返回:** 包含`tensor`塊的元組,分布在給定的`devices`上。
```
torch.cuda.comm.gather(tensors, dim=0, destination=None)
```
從多個GPU收集張量。
張量尺寸在不同于`dim`的所有維度上都應該匹配。
**參數:**
- **tensors** (*Iterable\[Tensor\]*) – 要收集的張量的迭代。
- **dim** (*int*) – 沿著此維度張量將被連接。
- **destination** (*int*, optional) – 輸出設備(-1表示CPU,默認值:當前設備)。
**返回:** 一個張量位于`destination`設備上,這是沿著`dim`連接`tensors`的結果。
## 流和事件
```
class torch.cuda.Stream
```
CUDA流的包裝。
**參數:**
- **device** (*int*, optional) – 分配流的設備。
- **priority** (*int*, optional) – 流的優先級。較低的數字代表較高的優先級。
> query()
檢查所有提交的工作是否已經完成。
**返回:** 一個布爾值,表示此流中的所有核心是否完成。
> record\_event(event=None)
記錄一個事件。
**參數:** **event** (*Event*, optional) – 要記錄的事件。如果沒有給出,將分配一個新的。 **返回:** 記錄的事件。
> synchronize()
等待此流中的所有核心完成。
> wait\_event(event)
將所有未來的工作提交到流等待事件。
**參數:** **event** (*Event*) – 等待的事件
> wait\_stream(stream)
與另一個流同步。
提交到此流的所有未來工作將等待直到所有核心在調用完成時提交給給定的流。
```
class torch.cuda.Event(enable_timing=False, blocking=False, interprocess=False, _handle=None)
```
CUDA事件的包裝。
**參數:**
- **enable\_timing** (*bool*) – 指示事件是否應該測量時間(默認值:False)
- **blocking** (*bool*) – 如果為true,`wait()`將被阻塞(默認值:False)
- **interprocess** (*bool*) – 如果為true,則可以在進程之間共享事件(默認值:False)
> elapsed\_time(end\_event)
返回事件記錄之前經過的時間。
> ipc\_handle()
返回此事件的IPC句柄。
> query()
檢查事件是否已被記錄。
**返回:** 一個布爾值,指示事件是否已被記錄。
> record(stream=None)
記錄給定流的事件。
> synchronize()
與事件同步。
> wait(stream=None)
使給定的流等待事件。
- PyTorch 中文文檔
- 主頁
- 自動求導機制
- CUDA語義
- 擴展PyTorch
- 多進程最佳實踐
- 序列化語義
- torch
- torch.Tensor
- torch.Storage
- torch.nn
- torch.nn.functional
- torch.autograd
- torch.optim
- torch.nn.init
- torch.multiprocessing
- torch.legacy
- torch.cuda
- torch.utils.ffi
- torch.utils.data
- torch.utils.model_zoo
- torchvision
- torchvision.datasets
- torchvision.models
- torchvision.transforms
- torchvision.utils
- 致謝