这个包增加了对CUDA张量类型的支持,它实现了与CPU张量相同的功能,但是它们利用gpu进行计算。它是惰性初始化的,所以您总是可以导入它,并使用is_available()来确定您的系统是否支持CUDA。CUDA semantics提供了更多关于使用CUDA的细节。
torch.cuda.current_blas_handle
()[source]
返回指向当前cuBLAS句柄的cublasHandle_t指针。
torch.cuda.current_device
()[source]
返回当前选定设备的索引。
torch.cuda.current_stream
(device=None)[source]
返回给定设备当前选定的流。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备当前选择的流,如果设备为None(默认),则由current_device()给出。
torch.cuda.default_stream
(device=None)[source]
返回给定设备的默认流。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的默认流,如果设备为None(默认),则由current_device()提供。
class torch.cuda.device
(device)[source]
更改所选设备的上下文管理器。
参数
device (torch.device or int) – 要选择的设备索引。如果这个参数是负整数或None,那么它就是no-op。
torch.cuda.device_count
()[source]
返回可用的gpu数量。
class torch.cuda.device_of
(obj)[source]
将当前设备更改为给定对象的设备的上下文管理器。您可以同时使用张量和存储作为参数。如果一个给定的对象没有分配在GPU上,这是一个no-op。
参数
obj (Tensor or Storage) – 在选定设备上分配的对象。
torch.cuda.empty_cache
()[source]
释放缓存分配器当前持有的所有未占用的缓存内存,以便这些内存可以在其他GPU应用程序中使用,并在nvidia-smi中可见。
注意
empty_cache()不会增加PyTorch可用的GPU内存。有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.get_device_capability
(device=None)[source]
获取设备的cuda功能。
参数
device (torch.device or int, optional) – 用于返回设备功能的设备。如果这个参数是一个负整数,那么这个函数就是no-op。如果设备为None(默认),则使用current_device()提供的当前设备。
返回值
主要和次要cuda功能的设备,返回类型。
torch.cuda.get_device_name
(device=None)[source]
获取设备的名称。
参数
device (torch.device or int, optional) – 用于返回名称的设备。如果这个参数是一个负整数,那么这个函数就是no-op。如果设备为None(默认),则使用current_device()提供的当前设备。
torch.cuda.init
()[source]
初始化PyTorch的CUDA状态。如果您通过PyTorch的C API与它进行交互,可能需要显式地调用这个函数,因为在初始化之前,CUDA功能的Python绑定不会这样做。普通用户不应该需要这样做,因为PyTorch的所有CUDA方法都会根据需要自动初始化CUDA状态。如果CUDA状态已经初始化,则不执行任何操作。
torch.cuda.ipc_collect
()[source]
Force在CUDA IPC释放GPU内存后收集GPU内存。
注意
检查是否有任何已发送的CUDA张量可以从内存中清除。如果没有活动计数器,则强制关闭用于引用计数的共享内存文件。当生成器进程停止主动发送张量并希望释放未使用的内存时,此函数非常有用。
torch.cuda.is_available
()[source]
返回一个bool,指示CUDA当前是否可用。
torch.cuda.max_memory_allocated
(device=None)[source]
返回给定设备张量占用的最大GPU内存(以字节为单位)。默认情况下,这将返回自该程序开始以来分配的内存峰值。reset_max_memory_assigned()可用于重置跟踪此指标的起始点。例如,这两个函数可以测量训练循环中每个迭代的分配内存使用量峰值。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.max_memory_cached
(device=None)[source]
返回缓存分配器为给定设备管理的最大GPU内存(以字节为单位)。默认情况下,这将返回自该程序开始以来的峰值缓存内存。reset_max_memory_cached()可用于重置跟踪此指标的起始点。例如,这两个函数可以测量训练循环中每个迭代的峰值缓存内存量。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.memory_allocated
(device=None)[source]
返回给定设备的张量占用的当前GPU内存(以字节为单位)。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
这可能比nvidia-smi中显示的要少,因为缓存分配器可以保存一些未使用的内存,并且需要在GPU上创建一些上下文。有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.memory_cached
(device=None)[source]
返回缓存分配器为给定设备管理的当前GPU内存(以字节为单位)。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.reset_max_memory_allocated
(device=None)[source]
重置跟踪给定设备的张量占用的最大GPU内存的起始点。有关详细信息,请参见max_memory_assigned()。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.reset_max_memory_cached
(device=None)[source]
重置跟踪由给定设备的缓存分配器管理的最大GPU内存的起始点。有关详细信息,请参见max_memory_cached()。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.set_device
(device)[source]
设置当前设备。这个功能的使用是不鼓励有利于设备。在大多数情况下,最好使用CUDA_VISIBLE_DEVICES环境变量。
参数
device (torch.device or int) – 选定的设备。如果这个参数是负数,这个函数就是no-op。
torch.cuda.stream
(stream)[source]
选择给定流的上下文管理器。在其上下文中排队的所有CUDA内核都将在选定的流上排队。
参数
stream (Stream) – selected stream. This manager is a no-op if it’s None
.
注意
流是种每设备。如果所选的流不在当前设备上,此函数还将更改当前设备以匹配流。
torch.cuda.synchronize
(device=None)[source]
等待CUDA设备上所有流中的所有内核完成。
参数
device (torch.device or int, optional) – 用于同步的设备。如果设备为None(默认),则使用current_device()提供的当前设备。
torch.cuda.get_rng_state
(device='cuda')[source]
以字节张量的形式返回指定GPU的随机数生成器状态。
参数
device (torch.device or int, optional) – 返回的RNG状态的设备。默认值:“cuda”(即torch.device('cuda'),当前cuda设备)。
警告
这个函数急切地初始化CUDA。
torch.cuda.get_rng_state_all
()[source]
返回一个字节张量元组,表示所有设备的随机数状态。
torch.cuda.set_rng_state
(new_state, device='cuda')[source]
设置指定GPU的随机数生成器状态。
Parameters
torch.cuda.set_rng_state_all
(new_states)[source]
设置所有设备的随机数生成器状态。
参数
new_state (tuple of torch.ByteTensor) – 每个设备所需的状态
torch.cuda.manual_seed
(seed)[source]
设置为当前GPU生成随机数的种子。如果CUDA不可用,调用这个函数是安全的;在这种情况下,它将被静静地忽略。
参数
seed (int) – The desired seed.
警告
如果您使用的是多gpu模型,这个函数不足以获得确定性。要为所有gpu播种,请使用manual_seed_all()。
torch.cuda.manual_seed_all
(seed)[source]
设置在所有gpu上生成随机数的种子。如果CUDA不可用,调用这个函数是安全的;在这种情况下,它将被静静地忽略。
参数
seed (int) – The desired seed.
torch.cuda.seed
()[source]
将生成随机数的种子设置为当前GPU的随机数。如果CUDA不可用,调用这个函数是安全的;在这种情况下,它将被静静地忽略。
警告
如果你使用的是多GPU模型,这个函数只会在一个GPU上初始化种子。要初始化所有gpu,请使用seed_all()。
torch.cuda.seed_all
()[source]
将生成随机数的种子设置为所有gpu上的随机数。如果CUDA不可用,调用这个函数是安全的;在这种情况下,它将被静静地忽略。
torch.cuda.initial_seed
()[source]
返回当前GPU的当前随机种子。
警告
torch.cuda.comm.broadcast
(tensor, devices)[source]
参数
返回值
一个包含张量副本的元组,放在与张量的指标相对应的设备上。
torch.cuda.comm.broadcast_coalesced
(tensors, devices, buffer_size=10485760)[source]
向指定的gpu广播序列张量。首先将小张量合并到缓冲区中,以减少同步的数量。
Parameters
返回值
一个包含张量副本的元组,放在与张量的指标相对应的设备上。
torch.cuda.comm.reduce_add
(inputs, destination=None)[source]
从多个gpu求和张量。所有输入都应该有匹配的形状。
参数
返回值
一个包含所有输入的元素和的张量,放在目标设备上。
torch.cuda.comm.scatter
(tensor, devices, chunk_sizes=None, dim=0, streams=None)[source]
在多个gpu上散射张量。
参数
返回值
包含张量块的元组,分布在给定的设备上。
torch.cuda.comm.gather
(tensors, dim=0, destination=None)[source]
从多个gpu收集张量。不同于dim的张量大小必须匹配。
参数
返回值
一个位于目标器件上的张量,它是沿dim将张量连接起来的结果。
class torch.cuda.Stream
[source]
包装一个CUDA流。CUDA流是一个线性执行序列,属于一个特定的设备,独立于其他流。详见CUDA语义。
参数
query
()[source]
检查提交的所有工作是否已经完成。
返回值
一个布尔值,指示该流中的所有内核是否已完成。
record_event
(event=None)[source]
记录一个事件。
参数
event (Event, optional) – event to record. If not given, a new one will be allocated.
返回值
记录事件。
synchronize
()[source]
等待流中的所有内核完成。
注意
这是一个cudaStreamSynchronize()的包装器:有关更多信息,请参见“CUDA文档”_。
wait_event
(event)[source]
让所有提交到流的未来工作等待一个事件。
参数
event (Event) – 要等待的事件。
注意
这是一个cudaStreamWaitEvent()的包装器:更多信息请参见“CUDA documentation”_。这个函数返回时不需要等待事件:只有未来的操作受到影响。
wait_stream
(stream)[source]
与另一个流同步。所有提交到此流的未来工作都将等待,直到调用时提交到给定流的所有内核都完成。
参数
stream (Stream) – a stream to synchronize.
注意
此函数返回时不需要在流中等待当前排队的内核:只有未来的操作受到影响。
class torch.cuda.Event
[source]
包装CUDA事件。CUDA事件是同步标记,可以用来监控设备的进程,准确地测量时间,并同步CUDA流。当事件首次被记录或导出到另一个进程时,底层的CUDA事件被惰性地初始化。创建后,只有同一设备上的流才能记录事件。然而,任何设备上的流都可以等待事件。
参数
elapsed_time
(end_event)[source]
返回事件被记录后和end_event被记录前经过的时间(以毫秒为单位)。
classmethod from_ipc_handle
(device, handle)[source]
从给定设备上的IPC句柄重构事件。
ipc_handle
()[source]
返回此事件的IPC句柄。如果尚未记录,事件将使用当前设备。
query
()[source]
检查事件当前捕获的所有工作是否已完成。
返回值
一个布尔值,指示当前由事件捕获的所有工作是否已完成。
record
(stream=None)[source]
在给定的流中记录事件。如果没有指定流,则使用torch.cuda.current_stream()。流的设备必须匹配事件的设备。
synchronize
()[source]
等待事件完成。直到完成此事件中当前捕获的所有工作。这将阻止CPU线程在事件完成之前继续执行。
注意
这是一个cudaEventSynchronize()的包装器:有关更多信息,请参见“CUDA documentation”_。
wait
(stream=None)[source]
使提交给给定流的所有未来工作等待此事件。如果没有指定流,那么使用torch.cuda.current_stream()。
torch.cuda.empty_cache
()[source]
释放缓存分配器当前持有的所有未占用的缓存内存,以便这些内存可以在其他GPU应用程序中使用,并在nvidia-smi中可见。
注意
empty_cache()不会增加PyTorch可用的GPU内存。有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.memory_allocated
(device=None)[source]
返回给定设备的张量占用的当前GPU内存(以字节为单位)。
参数
device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device()
, if device
is None
(default).
注意
这可能比nvidia-smi中显示的要少,因为缓存分配器可以保存一些未使用的内存,并且需要在GPU上创建一些上下文。有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.max_memory_allocated
(device=None)[source]
返回给定设备张量占用的最大GPU内存(以字节为单位)。默认情况下,这将返回自该程序开始以来分配的内存峰值。reset_max_memory_assigned()可用于重置跟踪此指标的起始点。例如,这两个函数可以测量训练循环中每个迭代的分配内存使用量峰值。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.reset_max_memory_allocated
(device=None)[source]
重置跟踪给定设备的张量占用的最大GPU内存的起始点。有关详细信息,请参见max_memory_assigned()。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.memory_cached
(device=None)[source]
返回缓存分配器为给定设备管理的当前GPU内存(以字节为单位)。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.max_memory_cached
(device=None)[source]
返回缓存分配器为给定设备管理的最大GPU内存(以字节为单位)。默认情况下,这将返回自该程序开始以来的峰值缓存内存。reset_max_memory_cached()可用于重置跟踪此指标的起始点。例如,这两个函数可以测量训练循环中每个迭代的峰值缓存内存量。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.reset_max_memory_cached
(device=None)[source]
重置跟踪由给定设备的缓存分配器管理的最大GPU内存的起始点。有关详细信息,请参见max_memory_cached()。
参数
device (torch.device or int, optional) – 选定的设备。返回当前设备的统计信息,由current_device()给出,如果设备为None(缺省值)。
注意
有关GPU内存管理的更多细节,请参见内存管理。
torch.cuda.nvtx.mark
(msg)[source]
描述某一时刻发生的瞬时事件。
参数
msg (string) – 与事件关联的ASCII消息。
torch.cuda.nvtx.range_push
(msg)[source]
将范围推到嵌套范围跨度的堆栈上。返回启动的范围的从零开始的深度。
参数
msg (string) – 与范围相关联的ASCII消息
torch.cuda.nvtx.range_pop
()[source]
从嵌套范围范围堆栈中弹出一个范围。返回结束的范围的从零开始的深度。