快捷方式

ExternalStream

class torch.cuda.ExternalStream(stream_ptr, device=None, **kwargs)[source][source]

外部分配的 CUDA 流的包装器。

此类用于包装在其他库中分配的流,以便于数据交换和多库交互。

注意

此类不管理流生命周期,用户有责任在使用此类时保持引用的流处于活动状态。

参数
  • stream_ptr (int) – cudaStream_t 值的整数表示形式。外部分配。

  • device (torch.deviceint, 可选) – 最初分配流的设备。如果设备指定不正确,则后续使用此流的启动可能会失败。

query()[source]

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

返回

一个布尔值,指示此流中的所有内核是否已完成。

返回类型

布尔型

record_event(event=None)[source]

记录事件。

参数

event (torch.cuda.Event, 可选) – 要记录的事件。如果未给出,将分配一个新的事件。

返回

已记录的事件。

synchronize()[source]

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

注意

这是 cudaStreamSynchronize() 的包装器:请参阅 CUDA Stream 文档 以获取更多信息。

wait_event(event)[source]

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

参数

event (torch.cuda.Event) – 要等待的事件。

注意

这是 cudaStreamWaitEvent() 的包装器:请参阅 CUDA Stream 文档 以获取更多信息。

此函数返回时不等待 event:仅影响未来的操作。

wait_stream(stream)[source]

与另一个流同步。

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

参数

stream (Stream) – 要同步的流。

注意

此函数返回时不等待 stream 中当前排队的内核:仅影响未来的操作。

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源