模块手册#

模块内容#

开物 PyTorch 插件公共 API。

class kaiwu.torch_plugin.BoltzmannMachine(num_nodes: int, quadratic_coef: FloatTensor | None = None, linear_bias: FloatTensor | None = None, device=None)[源代码]#

基类：AbstractBoltzmannMachine

玻尔兹曼机。

Args:

num_nodes (int): 模型中的节点总数。

quadratic_coef (torch.FloatTensor, optional): 二次项系数，

形状为 [num_nodes, num_nodes]

linear_bias (torch.FloatTensor, optional): 线性偏置，形状为 [num_nodes]

device (torch.device, optional): 构造张量所用的设备。若为 None，则使用 CPU。

clip_parameters(h_range, j_range) → None[源代码]#

原地裁剪线性和二次偏置权重。

Args:: h_range (tuple[float, float]): 二次项权重范围，例如 [-1, 1]。j_range (tuple[float, float]): 线性权重范围，例如 [-1, 1]。

condition_sample(sampler, s_visible, dtype=torch.float32) → Tensor[源代码]#

在给定部分节点的条件下从玻尔兹曼机中采样。

Args:

sampler (kaiwu.core.Optimizer): 用于采样的优化器。s_visible: 可见层状态。

Returns:

torch.Tensor: 从模型中采样得到的自旋: （形状由 sampler 和 sample_params 决定）。

forward(s_all: Tensor) → Tensor[源代码]#

计算哈密顿量。

Args:

s_all (torch.tensor): 形状为 (B, N) 的张量，其中 B 为批次大小，: N 为模型中的变量数。

Returns:

torch.tensor: 形状为 (B,) 的哈密顿量。

gibbs_sample(num_steps: int = 100, s_visible: Tensor | None = None, num_sample=None) → Tensor[源代码]#

从玻尔兹曼机中采样。

Args:

num_steps (int): 吉布斯采样步数。

s_visible (torch.Tensor, optional): 可见层状态，: 形状为 (B, num_visible)。若为 None，则随机初始化可见层。
num_sample (int, optional): 样本数。: 若为 None，则使用 s_visible 的批次大小。

hidden_bias(num_hidden: int) → Tensor[源代码]#

获取隐藏层偏置。

Args:: num_hidden (int): 隐藏节点数。

symmetrized_quadratic_coef()[源代码]#: 二次项系数

visible_bias(num_visible) → Tensor[源代码]#

获取可见层偏置。

Args:: num_visible (int): 可见节点数。

class kaiwu.torch_plugin.EnergyModel(bm_num_visible: int | None = None, bm_num_hidden: int | None = None, sampler: Any | None = None)[源代码]#

基类：Module

用于 QDiffusion 候选项评分的能量侧模型接口。

discretize_visible_state(visible_logits: Tensor) → Tensor[源代码]#: 将可见层 logits 转换为归一化的玻尔兹曼机可见层条件。

forward(noisy_tokens: Tensor, candidate_tokens: Tensor, attention_mask: Tensor) → Tensor | None[源代码]#: 通过 PyTorch 模块 API 执行条件能量评分。

get_last_stats() → dict[str, Tensor][源代码]#: 返回上一次评分调用的轻量级采样器诊断信息。

sample_hidden_state(visible_state: Tensor) → tuple[Tensor, Tensor][源代码]#: 为每个可见层赋值采样玻尔兹曼机隐藏层状态。

score_conditioned(noisy_tokens: Tensor, candidate_tokens: Tensor, attention_mask: Tensor) → Tensor | None[源代码]#: 对以噪声 token 为条件的候选项进行评分。

score_visible_logits(visible_logits: Tensor) → Tensor[源代码]#: 在条件玻尔兹曼机能量模型下对可见层 logits 进行评分。

class kaiwu.torch_plugin.QDiffusion(proposal_model: Module, energy_model: EnergyModel, token_spec: SequenceTokenSpec, config: QDiffusionConfig | None = None, dtype: dtype = torch.float32, device: device | str | None = None, freeze_proposal: bool = True)[源代码]#

基类：Module

Energy-guided discrete diffusion wrapper over generic sequence backbones. Initializes a QDiffusion model.

该类组合了两种骨干网络角色：

预测当前噪声状态 token logits 的提议模型
对候选重建结果重新排序的能量模型

该类同时提供面向训练的 API（如 :meth:objective）和面向解码的 API（如 :meth:initialize_state、:meth:step 和 :meth:generate）。

Args:

proposal_model: Backbone used to predict proposal logits.

energy_model: Energy-side model used to encode and score candidates.

token_spec: Special-token metadata required by the generator.

config: Optional generation/training configuration.

dtype: Floating point dtype tracked by the wrapper.

device: Optional target device. When omitted, infer from parameters.

freeze_proposal: Whether to freeze proposal model parameters.

energy(noisy_tokens: Tensor, candidate_tokens: Tensor, attention_mask: Tensor | None = None) → Tensor[源代码]#

对以噪声状态为条件的候选重建结果进行评分。

Args:

noisy_tokens: Noisy token tensor used as conditioning input.

candidate_tokens: Candidate clean token tensor to score.

attention_mask: Optional attention mask for the energy model.

Returns:

torch.Tensor: 形状为 [batch, 1] 的标量能量张量。

forward(noisy_tokens: Tensor, **kwargs: Any) → Tensor[源代码]#

在当前噪声状态上运行提议模型。

Args:: noisy_tokens: 当前噪声 token 张量。**kwargs: 转发给提议

模型的其他关键字参数。
Returns:: torch.Tensor: token 词表上的提议 logits。
Raises:: TypeError: 提议模型未实现 forward 时抛出。

generate(input_tokens: Tensor, *, max_steps: int = 500, partial_masks: Tensor | None = None, temperature: float = 1.0, return_state: bool = False) → Tensor | dict[str, Any][源代码]#

在核心类中运行完整的迭代解码循环。

Args:: input_tokens: 初始 token 张量。max_steps: 解码迭代次数。partial_masks: 固定位置的可选布尔掩码。temperature: 解码状态中的采样温度。return_state: 是否返回完整最终状态字典。
Returns:: torch.Tensor | dict[str, Any]: 最终 token 张量或完整解码状态。

get_non_special_symbol_mask(output_tokens: Tensor, partial_masks: Tensor | None = None) → Tensor[源代码]#

返回可编辑非特殊 token 位置的布尔掩码。

Args:: output_tokens: 待检查的 token 张量。partial_masks: 应保持

固定的位置的可选布尔掩码。
Returns:: torch.Tensor: 布尔掩码，其中 True 表示可编辑的非特殊位置。

initialize_state(input_tokens: Tensor, partial_masks: Tensor | None = None, max_steps: int = 500, temperature: float = 1.0) → dict[str, Any][源代码]#

为外部生成循环创建初始解码状态。

Args:: input_tokens: 初始 token 张量。partial_masks: 固定位置的可选布尔掩码。max_steps: 计划的解码迭代次数。temperature: 状态载荷中的采样温度。
Returns:: dict[str, Any]: 适用于重复调用 :meth:step 的可变状态字典。

objective(batch: dict[str, Tensor], weighting: str = 'constant') → dict[str, Tensor][源代码]#

构建外部循环使用的单步训练目标。

Args:: batch: 至少包含 batch["targets"] 的批次字典。weighting: 逐样本时间步加权模式。
Returns:: dict[str, torch.Tensor]: 包含提议 logits、监督掩码、损失权重和 EBM 目标项的字典。

proposal(noisy_tokens: Tensor, **kwargs: Any) → Tensor[源代码]#

:meth:forward 的语义别名，用于提议侧调用。

Args:: noisy_tokens: 当前噪声 token 张量。**kwargs: 转发给提议

模型的其他关键字参数。
Returns:: torch.Tensor: token 词表上的提议 logits。

step(state: dict[str, Any], partial_masks: Tensor | None = None) → dict[str, Any][源代码]#

执行一步去噪/重排序并返回更新后的状态。

Args:: state: 由 :meth:initialize_state 创建的当前解码状态。partial_masks: 固定位置的可选布尔掩码。
Returns:: dict[str, Any]: 一次迭代后更新的解码状态。

to(*args: Any, **kwargs: Any) → QDiffusion[源代码]#

移动模块并刷新缓存的设备和 dtype 元数据。

Args:: *args: 转发给 nn.Module.to 的位置参数。**kwargs: 转发给 nn.Module.to 的关键字参数。
Returns:: QDiffusion: 移动后的模块实例。

class kaiwu.torch_plugin.QDiffusionConfig(num_diffusion_timesteps: int = 500, use_coupled_sampling: bool = False, num_candidates: int = 1, proposal_temperature: float = 0.0, proposal_noise_scale: float = 1.0, energy_temperature: float = 1.0, disable_resample: bool = False, resample_ratio: float = 0.25, resample_top_p: float = 0.95, decoding_strategy: str = 'reparam-uncond-deterministic-linear')[源代码]#

基类：object

能量引导离散生成的配置。

Attributes:

num_diffusion_timesteps: 训练目标使用的离散加噪步数
。

use_coupled_sampling: Whether to use the coupled corruption variant.

num_candidates: Number of proposal candidates sampled at each decode step.

proposal_temperature: Temperature used for proposal-side sampling.

proposal_noise_scale: Gumbel noise scale used during proposal sampling.

energy_temperature: Temperature used when converting energies into
重排序权重时使用的温度。

disable_resample: Whether to disable repetition-collapse resampling.

resample_ratio: Frequency threshold that triggers resampling.

resample_top_p: Top-p cutoff used during resampling.

decoding_strategy: Skeptical-remasking strategy string.

decoding_strategy: str = 'reparam-uncond-deterministic-linear'#

disable_resample: bool = False#

energy_temperature: float = 1.0#

num_candidates: int = 1#

num_diffusion_timesteps: int = 500#

proposal_noise_scale: float = 1.0#

proposal_temperature: float = 0.0#

resample_ratio: float = 0.25#

resample_top_p: float = 0.95#

use_coupled_sampling: bool = False#

class kaiwu.torch_plugin.QVAE(input_dimension=None, activation_fct=None, config=None, **kwargs)[源代码]#

基类：AutoEncoderBase

Quantum Variational Autoencoder integrated into AutoEncoderBase framework.

Args:

input_dimension (int or list of int): Dimensionality of input features. activation_fct (callable, optional): Activation function for hidden layers. config (object): Configuration object containing:

num_latent_units (int)

loss_type (str): 'bernoulli' or 'mse'

dist_beta (float, default=10.0)

kl_beta (float, default=1e-6)

weight_decay (float, default=0.01)

sampler_type (str): Type of sampler for BM ('sa' or 'cim'). n_batches (int): Number of batches for conditional decoding (0 = no conditioning). bm (object, optional): Pre-created Boltzmann Machine. encoder (object, optional): Pre-created encoder. decoder (object, optional): Pre-created decoder. sampler (object, optional): Pre-created sampler. **kwargs: Additional keyword arguments for AutoEncoderBase.

bm_loss(q, bm_weight_decay=0.0)[源代码]#

Compute BM loss for updating BM parameters separately.

Args:

q (torch.Tensor): Encoder output logits (batch_size, latent_dim).: Must be detached to prevent gradients flowing to encoder.

bm_weight_decay (float, optional): L2 regularization coefficient for BM parameters.

Returns:

torch.Tensor: BM loss scalar.

create_networks()[源代码]#: Create encoder, decoder, BM and sampler. Subclasses must override.

energy(x, loss_type)[源代码]#: 计算 BM 能量（用于评估）

forward(x, batch_idx=None)[源代码]#

Forward pass through the QVAE.

Args:

x (torch.Tensor): Input tensor of shape (batch_size, input_dim). batch_idx (torch.Tensor, optional): Batch indices for conditional decoding.

Required if n_batches > 0.

Returns:

tuple: (recon_x, posterior, q, zeta)

recon_x: Reconstructed logits (batch_size, input_dim)
posterior: Posterior distribution object (MixtureGeneric)
q: Encoder output logits (batch_size, latent_dim)
zeta: Reparameterized latent sample (batch_size, latent_dim)

loss(x, recon_x, posterior)[源代码]#

Compute total loss (reconstruction + KL + weight decay).

Args:: x (torch.Tensor): Input tensor (batch_size, input_dim). recon_x (torch.Tensor): Reconstructed logits (batch_size, input_dim). posterior (MixtureGeneric): Posterior distribution object. q (torch.Tensor): Encoder logits (batch_size, latent_dim). zeta (torch.Tensor): Latent sample (batch_size, latent_dim).
Returns:: torch.Tensor: Total loss scalar.
Raises:: ValueError: If loss_type is not supported.

posterior(q_logits, beta)[源代码]#

Compute posterior distribution and reparameterized sample.

Args:

q_logits (torch.Tensor): Encoder output logits (batch_size, latent_dim). beta (float): Mixture parameter for MixtureGeneric.

Returns:

tuple: (posterior_dist, zeta)

posterior_dist: MixtureGeneric object
zeta: Reparameterized sample (batch_size, latent_dim)

set_train_bias(mean)[源代码]#: Compute train bias from dataset mean for Bernoulli reconstruction.

class kaiwu.torch_plugin.RestrictedBoltzmannMachine(num_visible: int, num_hidden: int, quadratic_coef: FloatTensor | None = None, linear_bias: FloatTensor | None = None, device=None)[源代码]#

基类：AbstractBoltzmannMachine

创建受限玻尔兹曼机。

Args:

num_visible (int): 模型中的可见节点数。

num_hidden (int): 模型中的隐藏节点数。

quadratic_coef (torch.FloatTensor, optional): 二次项系数，: 形状为 [num_visible, num_hidden]

linear_bias (torch.FloatTensor, optional): 线性偏置，形状为 [num_hidden]

device (torch.device, optional): 构造张量所用的设备。

clip_parameters(h_range, j_range) → None[源代码]#

原地裁剪线性和二次偏置权重。

Args:: h_range (tuple[float, float]): 二次项权重范围，例如 [-1, 1]。j_range (tuple[float, float]): 线性权重范围，例如 [-1, 1]。

forward(s_all: Tensor) → Tensor[源代码]#

计算哈密顿量。

Args:

s_all (torch.tensor): 形状为 (B, N) 的张量，其中 B 为批次大小，: N 为模型中的变量数。

Returns:

torch.tensor: 形状为 (B,) 的哈密顿量。

get_hidden(s_visible: Tensor, requires_grad: bool = False, bernoulli: bool = False) → Tensor[源代码]#

将可见层自旋传播到隐藏层。

Args:: s_visible: 可见层张量。requires_grad: 是否允许梯度反向传播。

get_visible(s_hidden: Tensor, bernoulli: bool = False) → Tensor[源代码]#: 将隐藏层自旋传播到可见层。

property hidden_bias: Tensor#: 返回隐藏层偏置。

property visible_bias: Tensor#: 返回可见层偏置。

class kaiwu.torch_plugin.UnsupervisedDBN(hidden_layers_structure=None)[源代码]#

基类：Module

通用的无监督深度信念网络（DBN）架构。

该模型由受限玻尔兹曼机（RBM）堆叠而成。

Args:

hidden_layers_structure (list, optional): 整数列表，: 表示每层的隐藏单元数。默认为 [100, 100]。

create_rbm_layer(input_dim)[源代码]#

为 DBN 创建各个 RBM 层。

Args:: input_dim (int): 输入数据维度（可见单元数）。
Returns:: UnsupervisedDBN: 已创建 RBM 层的实例本身。

forward(data_in)[源代码]#

执行前向传播以转换输入数据。

Args:: data_in (numpy.ndarray): 输入数据。
Returns:: numpy.ndarray: 经过所有 RBM 层后的转换数据。
Raises:: ValueError: 模型尚未构建或训练时抛出。

get_rbm_layer(index)[源代码]#

获取指定索引处的 RBM 层。

Args:: index (int): RBM 层索引。
Returns:: RestrictedBoltzmannMachine or None: 找到时返回 RBM 层，否则返回 None。

mark_as_trained()[源代码]#

将模型标记为已训练。

Returns:: UnsupervisedDBN: 实例本身。

property num_layers#

返回 RBM 层数。

Returns:: int: 层数。

property output_dim#

返回 DBN 的输出维度。

Returns:: int: 最终隐藏层的维度。

reconstruct(data_in, layer_index=0)[源代码]#

从指定 RBM 层重建输入。

Args:

data_in (numpy.ndarray): 待重建的输入数据。

layer_index (int, optional): 用于重建的 RBM 层索引。: 默认为 0。

Returns:

numpy.ndarray: 重建后的数据。

Raises:

ValueError: 模型没有 RBM 层或层索引超出范围时抛出。

static reconstruct_with_rbm(rbm, data_in, device=None)[源代码]#

使用单个 RBM 重建数据。

Args:

rbm (RestrictedBoltzmannMachine): 已训练的 RBM 模型。

data_in (numpy.ndarray): 输入数据。

device (torch.device, optional): 执行计算所用的设备。: 若为 None，则使用 RBM 的设备。默认为 None。

Returns:

tuple[numpy.ndarray, numpy.ndarray]: 包含以下内容的元组：

重建后的可见层数据。
每个样本的重建误差。

transform(data_in)[源代码]#

兼容 sklearn 的 transform 方法。

Args:: data_in (numpy.ndarray): 输入数据。
Returns:: numpy.ndarray: 转换后的数据。

kaiwu.torch_plugin.disable_usage_stats() → None[源代码]#: 禁用 kpp 使用情况统计。禁用后不会上报任何数据。

kaiwu.torch_plugin.enable_usage_stats() → None[源代码]#: 启用 kpp 使用情况统计（默认启用，通常无需调用）。

kaiwu.torch_plugin.is_usage_stats_enabled() → bool[源代码]#: 返回当前是否已启用使用情况统计。