spikingjelly.activation_based.cuda_utils package

Module contents

spikingjelly.activation_based.cuda_utils.cpu_timer(f: Callable, *args, **kwargs)[源代码]

计算在CPU上执行 f(*args, **kwargs) 所需的时间

参数:

f (Callable) – 函数

返回:

用时，单位是毫秒

返回类型:

float

Returns the used time for calling f(*args, **kwargs) in CPU

参数:

f (Callable) – a function

返回:

used time in milliseconds

返回类型:

float

spikingjelly.activation_based.cuda_utils.cuda_timer(device: Union[device, int], f: Callable, *args, **kwargs)[源代码]

计算在CUDA上执行 f(*args, **kwargs) 所需的时间

参数:

  • device (Union[torch.device, int]) – f 运行的CUDA设备

  • f (Callable) – 函数

返回:

用时，单位是毫秒

返回类型:

float

Returns the used time for calling f(*args, **kwargs) in CUDA

参数:

  • device (Union[torch.device, int]) – on which cuda device that f is running

  • f (Callable) – a function

返回:

used time in milliseconds

返回类型:

float

spikingjelly.activation_based.cuda_utils.cal_fun_t(n: int, device: Union[str, device, int], f: Callable, *args, **kwargs)[源代码]

测量在 device 上执行 n 次 f(*args, **kwargs) 的平均用时

备注

当 n > 1 时，实际上会执行 2n 次，然后返回后 n 次的平均用时，以减小误差。

参数:

  • n (int) – 重复的次数

  • device (Union[str, torch.device, int]) – f 执行的设备，可以为 ‘cpu’ 或CUDA设备

  • f (Callable) – 函数

返回:

用时，单位是毫秒

返回类型:

float

Returns the used time averaged by calling f(*args, **kwargs) over n times

Note

If n > 1, this function will call f for 2n times and return the average used time by the last n times to reduce the measure error.

参数:

  • n (int) – repeat times

  • device (Union[str, torch.device, int]) – on which cuda device that f is running. It can be ‘cpu’ or a cuda deivce

  • f (Callable) – function

返回:

used time in milliseconds

返回类型:

float

spikingjelly.activation_based.cuda_utils.cal_blocks(numel: int, threads: int = -1)[源代码]
参数:

  • numel (int) – 并行执行的CUDA内核的数量

  • threads (int) – 每个cuda block中threads的数量，默认为-1，表示使用 configure.cuda_threads

返回:

blocks的数量

返回类型:

int

此函数返回 blocks的数量，用来按照 kernel((blocks,), (configure.cuda_threads,), ...) 调用 cupy.RawKernel

参数:

  • numel (int) – the number of parallel CUDA kernels

  • threads (int) – the number of threads in each cuda block. The defaule value is -1, indicating to use configure.cuda_threads

返回:

the number of blocks

返回类型:

int

Returns the number of blocks to call cupy.RawKernel by kernel((blocks,), (threads,), ...)

spikingjelly.activation_based.cuda_utils.get_contiguous(*args)[源代码]

将 *args 中所有的 torch.Tensor 或 cupy.ndarray 进行连续化。

备注

连续化的操作无法in-place，因此本函数返回一个新的list。

返回:

一个元素全部为连续的 torch.Tensor 或 cupy.ndarray 的 list

返回类型:

list

返回:

a list that contains the contiguous torch.Tensor or cupy.ndarray

返回类型:

list

Makes torch.Tensor or cupy.ndarray in *args to be contiguous

Note

The making contiguous operation can not be done in-place. Hence, this function will return a new list.

spikingjelly.activation_based.cuda_utils.wrap_args_to_raw_kernel(device: int, *args)[源代码]
参数:

device (int) – raw kernel运行的CUDA设备

返回:

一个包含用来调用 cupy.RawKernel 的 tuple

返回类型:

tuple

此函数可以包装 torch.Tensor 和 cupy.ndarray 并将其作为 cupy.RawKernel.__call__ 的 args

参数:

device (int) – on which CUDA device the raw kernel will run

返回:

a tuple that contains args to call cupy.RawKernel

返回类型:

tuple

This function can wrap torch.Tensor or cupy.ndarray to args in cupy.RawKernel.__call__

class spikingjelly.activation_based.cuda_utils.DeviceEnvironment(device: int)[源代码]

基类：object

这个模块可以被用作在指定的 device 上执行CuPy函数的上下文，用来避免 torch.cuda.current_device() 被CuPy意外改变( https://github.com/cupy/cupy/issues/6569 )。

代码示例：

with DeviceEnvironment(device):
    kernel((blocks,), (configure.cuda_threads,), ...)
参数:

device (int) – the CUDA device

This module is used as a context to make CuPy use the specific device, and avoids torch.cuda.current_device() is changed by CuPy ( https://github.com/cupy/cupy/issues/6569 ).

Codes example:

with DeviceEnvironment(device):
    kernel((blocks,), (configure.cuda_threads,), ...)