14.md

# torch.cuda

*   [交流集](https://ptorch.com/docs/1/torch-cuda#communication-collectives)
*   [流和事件](https://ptorch.com/docs/1/torch-cuda#streams-and-events)
*   [NVIDIA 工具扩展（NVTX）](https://ptorch.com/docs/1/torch-cuda#nvidia-tools-extension-nvtx)

* * *

该包增加了对 CUDA 张量类型的支持，实现了与 CPU 张量相同的功能，但使用 GPU 进行计算。

它是延迟的初始化，所以你可以随时导入它，并使用`is_available()`来确定系统是否支持 CUDA。

[CUDA 语义](http://pytorch.org/docs/master/notes/cuda.html#cuda-semantics)有关于使用 CUDA 的更多细节。

```py
torch.cuda.current_blas_handle() 
```

返回指向当前 cuBLAS 句柄的 cublasHandle_t 指针

```py
torch.cuda.current_device() 
```

返回当前所选设备的索引。

```py
torch.cuda.current_stream() 
```

返回当前选定的`Stream`

```py
class torch.cuda.device(idx) 
```

更改所选设备的上下文管理器。

参数：

*   idx(int) – 设备索引选择。如果这个参数是负的，则是无效操作。

```py
torch.cuda.device_count() 
```

返回可用的 GPU 数量。

```py
class torch.cuda.device_of(obj) 
```

将当前设备更改为给定对象的上下文管理器。

可以使用张量和存储作为参数。如果给定的对象不是在 GPU 上分配的，则是无效操作。

参数：

*   obj (Tensor 或者 Storage) – 在选定设备上分配的对象。

```py
torch.cuda.is_available() 
```

返回 bool 值，指示当前 CUDA 是否可用。

```py
torch.cuda.set_device(device) 
```

设置当前设备。

不鼓励使用此功能函数。在大多数情况下，最好使用`CUDA_VISIBLE_DEVICES`环境变量。

参数：

*   device(int) - 选择的设备。如果此参数为负，则此函数是无操作的。

```py
torch.cuda.stream(stream) 
```

选择给定流的上下文管理器。

在其上下文中排队的所有 CUDA 核心将在所选流上排列。

参数：

*   stream(Stream) – 选择的流。如果为`None`，则这个管理器是无效的。

```py
torch.cuda.synchronize() 
```

等待当前设备上所有流中的所有内核完成。

### 交流集

```py
torch.cuda.comm.broadcast(tensor, devices) 
```

向一些 GPU 广播张量。

参数：

*   tensor (Tensor) – 广播的张量
*   devices (Iterable) – 可以广播的设备的迭代。注意，它的设备形式为（src，dst1，dst2，...），其第一个元素是广播来源的设备。

返回： 包含张量副本的元组，放置在对应于索引的设备上。

```py
torch.cuda.comm.reduce_add(inputs, destination=None) 
```

将来自多个 GPU 的张量相加。

所有输入应具有匹配的形状。

参数：

*   inputs (Iterable[Tensor]) – 要相加张量的迭代
*   destination (int, 可选) – 将放置输出的设备（默认值：当前设备）。

返回： 包含放置在`destination`设备上的所有输入的元素总和的张量。

```py
torch.cuda.comm.scatter(tensor, devices, chunk_sizes=None, dim=0, streams=None) 
```

跨多个 GPU 分散张量。

参数：

*   tensor (Tensor) – 要分散的张量
*   devices (Iterable[int]) – 可迭代的 int，指定张量应分散在哪些设备中。
*   chunk_sizes (Iterable[int], 可选) – 要放置在每个设备上的块大小。它应该匹配`devices`的长度且总和为`tensor.size(dim)`。 如果没有指定，张量将被分成相等的块。
*   dim (int, 可选) – 将张量块化的一个维度。

返回： 包含`tensor`块的元组，传播给`devices`。

```py
torch.cuda.comm.gather(tensors, dim=0, destination=None) 
```

从多个 GPU 收集张量。

所有张量的测量尺寸与 dim 不一致。

参数：

*   tensors (Iterable[Tensor]) – 要收集的张量的迭代。
*   dim (int) – 张量沿着此维度来连接。
*   destination (int, 可选) – 输出设备（-1 表示 CPU，默认值：当前设备）。

返回： 一个张量位于`destination`设备上，这是沿着`dim`连接`tensors`的结果。

## 流和事件

```py
class torch.cuda.Stream 
```

CUDA 流的包装。

参数：

*   device (int, 可选) – 分配流的设备。
*   priority (int, 可选) – 流的优先级。数字越小优先级越高。

    > query()

    检查所有提交的工作是否已经完成。

    返回： 一个布尔值，表示此流中的所有核心是否完成。

    > record_event(event=None)

    记录事件。

    参数：

    *   event (Event, 可选) – 要记录的事件。如果未定义，将分配一个新的。 返回： 记录的事件。

    > synchronize()

    等待此流中的所有核心完成。

    > wait_event(event)

    将所有未来的工作提交到流等待事件。

    参数： event (Event) – 等待的事件

    > wait_stream(stream)

    与另一个流同步。

    提交到此流的所有未来工作将等待直到所有核心在调用完成时提交给给定的流。

```py
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)

    使给定的流等待事件。

    ## NVIDIA 工具扩展（NVTX）

    ```py
    torch.cuda.nvtx.mark(msg) 
    ```

    描述在某个时刻发生的瞬时事件。 参数：

    *   msg（string） - 与事件关联的 ASCII 消息。

    ```py
    torch.cuda.nvtx.range_push(msg) 
    ```

    将范围推送到嵌套范围跨度的堆栈。返回起始范围的基于 zero-based 的深度。 参数：

    *   msg（string） - 与范围关联的 ASCII 消息

        ```py
        torch.cuda.nvtx.range_pop() 
        ```

        从一堆嵌套的范围范围内弹出一个范围。返回结束的范围的基于 zero-based 的深度。

### 译者署名

| 用户名 | 头像 | 职能 | 签名 |
| --- | --- | --- | --- |
| [Song](https://ptorch.com) | ![](img/2018033000352689884.jpeg) | 翻译 | 人生总要追求点什么 |