spikingjelly.activation_based.memopt package#

本子包提供了减少 spikingjelly.activation_based 模型训练显存开销的工具。详情请参阅我们在 ICLR 2026 上发表的论文 Towards Lossless Memory-efficient Training of Spiking Neural Networks via Gradient Checkpointing and Spike Compression 以及源代码仓库。

This package provides tools for reducing training memory consumption of spikingjelly.activation_based models. See our ICLR 2026 paper Towards Lossless Memory-efficient Training of Spiking Neural Networks via Gradient Checkpointing and Spike Compression and source code repository for details.

Optimization Pipeline#

基于梯度检查点和脉冲压缩的深度SNN训练显存自动优化工具。

Automatic memory optimization pipeline for deep SNN training based on gradient checkpointing and spike compression.

`memory_optimization`	The main API. Perform memory optimization on a model.
`MemOptSummary`	Structured summary returned by `memory_optimization(..., return_summary=True)`.
`MEMOPT_PROFILES`	Supported high-level memopt presets.
`resolve_device`	Get the device of the current process.
`apply_gc`	Apply GC to a submodule.
`get_module_and_parent`	Get a module and its parent module given its path.

用于实现带输入压缩的梯度检查点 (GC) 的工具。

Tools for implementing gradient checkpointing (GC) with input compression.

`in_gc_1st_forward`	Whether in the first forward pass of GC.
`query_autocast`	Query autocast information.
`input_compressed_gc`	Wrap a function with GC and input compression.
`GCContainer`	Module container representing a GC segment.
`TCGCContainer`	Module container representing a temporally chunked GC segment.

将浮点数表示的脉冲张量转换为更紧凑的表示形式的压缩器。

Compressors that convert spike tensors represented in floating-point numbers into more compact representations.

`BaseSpikeCompressor`	Base class for spike compressors.
`NullSpikeCompressor`	Do not perform any compression/decompression.
`BooleanSpikeCompressor`	Convert spike tensors to/from boolean tensors.
`Uint8SpikeCompressor`	Convert spike tensors to/from `uint8` tensors.
`BitSpikeCompressor`	Converts spike tensors to/from bit representations.
`SparseSpikeCompressor`	Convert spike tensors to/from sparse representations.