模块 basalt.torch¶

39 个公开类型

PyTorch 风格 eager API：动态张量、Module 组合、原生 autograd、优化器与损失。

类型概览

final class Activation<L extends Layout> extends Module<L, L>¶

逐元素激活函数模块，保持输入布局

forward 按 kind 分派到 Tensor 上的一元算子：RELU、GELU、SILU、SIGMOID、TANH 直接转调同名方法，LEAKY_RELU 与 ELU 额外把 alpha 传给对应算子。输出与输入同形状，不做任何维度变换，也不改变布局标签。

模块不持有参数，parameters() 恒返回空列表，因此不会进入优化器的参数集；类型参数 L 的作用是让激活在 Sequential 中保持前后布局一致。

注意

kind 不做 null 校验，构造时传入 null 会在 forward 的 switch 分派处抛 NullPointerException。

注意

激活由原生内核逐元素计算，输出是新分配的张量，梯度按元素导数回传。

参见：basalt.torch.Module, basalt.torch.Sequential, basalt.torch.Activation.Kind

构造器

Activation(Kind kind)¶

按激活种类构造无参模块

注意

构造器只保存枚举值，不分配任何张量。

参数

• kind (Kind) – 激活种类，取值来自 Activation.Kind；不做 null 校验

Activation(Kind kind, float alpha)¶

按激活种类与附加系数构造

注意

单参构造器按 kind 取默认值：LEAKY_RELU 为 0.01f，其余种类为 1.0f。

注意

构造器只保存枚举值与系数，不分配任何张量。

参数

• kind (Kind) – 激活种类，取值来自 Activation.Kind；不做 null 校验

• alpha (float) – 附加系数：LEAKY_RELU 的负半轴斜率、ELU 的负半轴饱和值；其余种类不读取该值，也不校验范围

方法

kind() → Kind¶

返回激活种类

返回

构造时传入的 Kind 枚举值

alpha() → float¶

返回附加系数

返回

构造时传入的 alpha；RELU、GELU、SILU、SIGMOID、TANH 分支不使用该值

forward(Tensor input) → Tensor¶

对输入逐元素应用激活

注意

训练与推理行为相同，切换 train()/eval() 不影响结果。

参数

• input (Tensor) – 任意秩、任意布局的张量；形状与元素个数不被改变

返回

与 input 同形状的新张量

parameters() → List<Tensor>¶

返回空参数列表

返回

恒为空的不可变列表

final class AdamW extends Optimizer¶

解耦权重衰减与偏差校正的 Adam 优化器

每步执行 p -= lr * (m̂ / (√v̂ + ε) + wd * p)：m̂、v̂ 是经偏差校正的一阶与二阶矩估计，权重衰减以加性项直接作用于参数，不进入矩估计，因此衰减强度与自适应步长相互独立。

一阶、二阶矩缓冲由原生上下文按张量地址保存，首次 step() 时按参数长度惰性分配；参数长度变化会重置缓冲。单参构造器采用 beta1=0.9、beta2=0.999、epsilon=1e-8、weightDecay=0.01。

注意

step() 先自增内部计数器再更新参数，计数器从 1 开始，正好是偏差校正所需的幂次；stepCount() 返回该计数。

注意

构造器不校验 beta1、beta2、epsilon、weightDecay 的取值，传入非有限值会直接污染参数。

注意

同一参数不应同时交给两个优化器：矩缓冲按张量地址存放在上下文中，后创建者会覆盖前者。

参见：basalt.torch.Optimizer, basalt.torch.SGD, basalt.torch.Tensor#adam(float, float, float, float, float, long)

构造器

AdamW(List<Tensor> parameters, float learningRate)¶

以默认超参数构造 AdamW

等价于 beta1=0.9、beta2=0.999、epsilon=1e-8、weightDecay=0.01 的完整构造器。

注意

列表在构造时被复制，之后修改原列表不影响优化器。

参数

• parameters (List<Tensor>) – 待优化的可训练张量列表，必须非空

• learningRate (float) – 步长，必须有限且大于 0

抛出

• IllegalArgumentException – 参数列表为空，或学习率非正、非有限

AdamW(List<Tensor> parameters, float learningRate, float beta1, float beta2, float epsilon, float weightDecay)¶

以显式超参数构造 AdamW

注意

beta1、beta2、epsilon、weightDecay 不做取值校验。

参数

• parameters (List<Tensor>) – 待优化的可训练张量列表，必须非空

• learningRate (float) – 步长，必须有限且大于 0

• beta1 (float) – 一阶矩的指数衰减率

• beta2 (float) – 二阶矩的指数衰减率

• epsilon (float) – 分母稳定项，避免二阶矩为零时除零

• weightDecay (float) – 解耦权重衰减系数，直接加在参数上

抛出

• IllegalArgumentException – 参数列表为空，或学习率非正、非有限

方法

step()¶

执行一步更新，并把内部步数加一

先自增步数计数器，再对每个参数调用 Tensor.adam，用同一计数器做偏差校正。

注意

不清零梯度，也不清计算图；梯度为零时更新量为零但仍会推进计数器。

stepCount() → long¶

返回已执行的更新步数

返回

从 0 开始、每次 step() 加一的计数

setStepCount(long step)¶

覆盖内部步数计数器

注意

只改计数器，不触碰参数，也不重置原生上下文里的一阶、二阶矩缓冲。

注意

恢复后首次 step() 会先自增，再以新的计数器计算 1 - beta^t。

参数

• step (long) – 新的步数，必须大于等于 0；用于从检查点续训时对齐偏差校正的幂次

抛出

• IllegalArgumentException – step 为负

final class BCEWithLogitsLoss implements Loss¶

数值稳定的带 logits 二分类交叉熵

loss(prediction, target) 直接转调 Tensor.bceWithLogitsLoss：prediction 是未过 sigmoid 的 logits，target 与 prediction 形状完全相同、元素取 0 或 1，返回形状 [1] 的标量，按元素总数取平均。

前向用 max(x, 0) - x·y + log1p(exp(-|x|)) 的等价形式逐元素累加，避免 exp 在正半轴溢出；反向不沿 sigmoid 的自动微分链回传，而是把 (sigmoid(x) - y)/n 作为常量系数构造一次 reshape 与 matMul，再叠加常数偏移凑出标量值，因此对 prediction 的一阶梯度精确。

示例

try (Basalt jsl = Basalt.create()) {
    BCEWithLogitsLoss criterion = new BCEWithLogitsLoss();
    Tensor logits = basalt.randn(4, 1);
    Tensor target = basalt.zeros(4, 1);   // 二分类标签，元素取 0 或 1
    Tensor loss = criterion.loss(logits, target);
    loss.backward();
}

注意

prediction 与 target 形状必须逐维相同，否则抛 IllegalArgumentException；不广播。

注意

target 只被读取原始数值，不参与求导，把它设为 requiresGrad 也不会得到梯度。

注意

返回张量在 prediction 上是线性的替代图，一阶梯度正确，但二阶导数不等于真实 BCE 的 Hessian。

注意

无状态、无可训练参数，可复用同一实例。

参见：basalt.torch.Loss, basalt.torch.CrossEntropyLoss, basalt.torch.Tensor#bceWithLogitsLoss(Tensor)

方法

loss(Tensor prediction, Tensor target) → Tensor¶

计算带 logits 的二分类交叉熵

注意

梯度按 (sigmoid(prediction) - target) / n 附着在 prediction 上，n 为元素总数而非批大小。

参数

• prediction (Tensor) – 任意形状的 logits 张量，与 target 逐维同形；内部会取 sigmoid，不要预先激活

• target (Tensor) – 与 prediction 同形状的标签张量，元素取 0 或 1；仅读取数值，不参与求导

返回

形状 [1] 的标量损失张量，值为全体元素的平均，可继续调用 backward()

抛出

• IllegalArgumentException – prediction 与 target 形状不一致

final class BatchNorm2d extends Module<BCHW, BCHW>¶

NCHW 上的批归一化，含滑动统计量

训练模式下按通道统计 N×H×W 上的均值与方差，用批内方差归一化后做仿射变换，同时按 running = (1 - momentum)·running + momentum·batch 的方向更新滑动均值与滑动方差，即 momentum 是新批统计量的权重。

推理模式不再统计批内数据，直接用滑动均值与滑动方差归一化。写入滑动方差的是无偏估计（spatial 大于 1 时乘 spatial / (spatial - 1)），而归一化本身使用有偏方差。

gamma 初始化为 1、beta 初始化为 0，两者是可训练参数；running_mean 初始化为 0、running_var 初始化为 1，它们是不可训练状态，仅在训练前向中被原地改写。

注意

默认 epsilon 为 1e-5、momentum 为 0.1，来自单参构造器；构造器不校验这两个值的范围。

注意

滑动统计量是模块状态而非参数：parameters() 只返回 gamma 与 beta，namedParameters() 额外给出 running_mean 与 running_var。

注意

在从未训练过的模块上直接 eval()，归一化使用均值 0、方差 1，输出近似为 gamma·x + beta。

参见：basalt.torch.LayerNorm, basalt.torch.Module#train(), basalt.torch.Tensor#batchNorm2d

构造器

BatchNorm2d(NativeContext context, int channels)¶

以 epsilon=1e-5、momentum=0.1 构造

注意

等价于调用四参构造器并传入 1e-5f 与 0.1f。

参数

• context (NativeContext) – 持有原生上下文，参数与状态张量都在其中分配

• channels (int) – 通道数 C，即 gamma、beta 与滑动统计量的长度，必须为正

抛出

• IllegalArgumentException – channels 小于等于 0

BatchNorm2d(NativeContext context, int channels, float epsilon, float momentum)¶

显式指定数值稳定项与滑动动量

注意

momentum 越界不会报错，但会让滑动统计量发散或完全停滞。

参数

• context (NativeContext) – 持有原生上下文，参数与状态张量都在其中分配

• channels (int) – 通道数 C，必须为正

• epsilon (float) – 加在方差上的数值稳定项，避免除零；不校验正负

• momentum (float) – 滑动统计量更新时新批统计量的权重，取 0 到 1 时表现为指数滑动平均；不校验范围

抛出

• IllegalArgumentException – channels 小于等于 0

方法

forward(Tensor input) → Tensor¶

按当前训练或推理模式做批归一化

注意

训练模式会原地改写 running_mean 与 running_var；推理模式只读。

注意

批内统计覆盖整个 N×H×W，batch 为 1 时方差可能接近 0，此时 epsilon 决定数值稳定性。

参数

• input (Tensor) – 形状 [N, C, H, W] 的张量，C 需与构造时的 channels 一致

返回

[N, C, H, W]，与输入同形状的新张量

抛出

• IllegalArgumentException – 输入不是四维张量

parameters() → List<Tensor>¶

返回可训练的 gamma 与 beta

返回

长度为 2 的列表，顺序为 gamma、beta

epsilon() → float¶

返回数值稳定项

返回

构造时传入的 epsilon，单参构造器下为 1e-5

gamma() → Tensor¶

返回缩放参数

返回

形状 [channels] 的可训练张量，初值全 1

beta() → Tensor¶

返回平移参数

返回

形状 [channels] 的可训练张量，初值全 0

runningMean() → Tensor¶

返回滑动均值状态

返回

形状 [channels] 的非可训练张量，初值全 0，训练前向中原地更新

runningVar() → Tensor¶

返回滑动方差状态

返回

形状 [channels] 的非可训练张量，初值全 1，训练前向中按无偏估计原地更新

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回含滑动统计量的命名状态

注意

stateDict 因此包含推理所需的统计量，loadStateDict 可完整恢复 eval 行为。

返回

gamma、beta、running_mean、running_var 四项，顺序固定

close()¶

释放参数与统计量张量

注意

关闭后模块不可再用；四个张量均由本模块独占。

final class Checkpoint¶

训练检查点的保存与恢复

save 把模块的 stateDict、epoch、loss、优化器学习率与 AdamW 步数写进单个二进制文件：魔数 0x42534C4B（“BasaltK”）、版本 1，随后是 epoch、loss、学习率、步数，再逐项写入名称长度 + UTF-8 名称 + 元素个数 + 以 double 存储的数值。

restore 先 load 再把参数表灌回模块：文件中学习率大于 0 时写回优化器，优化器是 AdamW 时同时恢复步数计数器，返回的 Checkpoint 提供 epoch 与 loss，供续训时对齐轮次与日志。

AdamW 的一阶、二阶矩缓冲保存在原生上下文里、按张量句柄索引，不进入检查点；恢复后矩缓冲从零开始累积，偏差校正按恢复出的步数计算，因此续训初期的更新幅度与不间断训练并不相同。

示例

try (Basalt jsl = Basalt.create()) {
    Linear model = new Linear(basalt.context(), 16, 4);
    AdamW optimizer = new AdamW(model.parameters(), 1e-3f);
    // ... 训练若干轮
    Checkpoint.save(Path.of("ckpt.bin"), model, optimizer, 3, 0.42f);
    Checkpoint ckpt = Checkpoint.restore(Path.of("ckpt.bin"), model, optimizer);
    int nextEpoch = ckpt.epoch() + 1;
}

注意

参数以 float 写盘、以 double 读回再截断为 float，往返无损；文件不含校验和，也不含模型结构。

注意

save 的 optimizer 允许为 null，此时学习率记 0、步数记 0，restore 侧因此不会覆盖优化器学习率。

注意

load 校验魔数与版本，读完后检查没有剩余字节，任一不符抛 ModelFormatException。

注意

恢复要求模块的 namedParameters 键集与文件一致，缺键抛 IllegalArgumentException；多余键被忽略。

参见：basalt.torch.Module#stateDict(), basalt.torch.AdamW#stepCount(), basalt.torch.CosineAnnealingLR

构造器

Checkpoint(int epoch, float loss, float learningRate, long optimizerStep, Map<String, float[]> stateDict)¶

以给定训练状态构造检查点对象

注意

映射被复制成不可修改视图，但其中的 float[] 是克隆副本，改动它们会改变本实例的取值。

注意

通常由 load 或 restore 产生，手工构造主要用于测试。

参数

• epoch (int) – 已完成的轮次序号，原样保存

• loss (float) – 该轮记录的损失值

• learningRate (float) – 保存时优化器的学习率，0 表示未记录

• optimizerStep (long) – AdamW 的累计步数，非 AdamW 优化器记 0

• stateDict (Map<String, float[]>) – 参数名到数值数组的映射，构造时逐项克隆

方法

save(Path target, Module<?, ?> module, Optimizer optimizer, int epoch, float loss)static¶

把模块状态与训练进度写入文件

注意

target 或 module 为 null 时抛 NullPointerException。

注意

写入是整文件覆盖，不做原子替换，也不保留旧版本。

参数

• target (Path) – 目标文件路径，父目录必须已存在；已存在时被整体覆盖

• module (Module<?, ?>) – 提供 stateDict 的模块，其 namedParameters 决定写入的键

• optimizer (Optimizer) – 可选，为 null 时学习率与步数都记 0；非 AdamW 优化器只记录学习率

• epoch (int) – 写入的轮次序号

• loss (float) – 写入的损失值

抛出

• UncheckedIOException – 写入文件失败，例如目录不存在或没有写权限

load(Path source) → Checkpointstatic¶

从文件读取检查点

注意

只做格式校验，不接触任何模块或优化器。

参数

• source (Path) – 已存在的检查点文件路径

返回

解析出的 Checkpoint，含 epoch、loss、学习率、优化器步数与参数表

抛出

• ModelFormatException – 魔数或版本不符、数据被截断、长度字段非法，或读取后仍有多余字节

• UncheckedIOException – 读取文件失败

restore(Path source, Module<?, ?> module, Optimizer optimizer) → Checkpointstatic¶

读取检查点并恢复到模块与优化器

注意

文件中学习率为 0 时跳过回写；非 AdamW 优化器不恢复步数。

注意

不恢复 AdamW 的矩缓冲；文件格式错误时抛出与 load 相同的异常。

参数

• source (Path) – 检查点文件路径

• module (Module<?, ?>) – 接收参数值的模块，其 namedParameters 的键必须覆盖文件中的全部条目

• optimizer (Optimizer) – 可选；非 null 时按文件内容回写学习率，AdamW 还会恢复步数计数器

返回

读取到的 Checkpoint，可从中取 epoch 与 loss

抛出

• IllegalArgumentException – 模块缺少文件中的某个状态键

epoch() → int¶

返回保存时的轮次序号

返回

检查点记录的 epoch，通常表示已完成的轮数

loss() → float¶

返回保存时记录的损失

返回

写入时的 float 损失值

learningRate() → float¶

返回保存时的学习率

返回

写入时的学习率；save 时 optimizer 为 null 则为 0

optimizerStep() → long¶

返回保存时的优化器步数

返回

AdamW 的累计步数；其他优化器与无优化器情形为 0

stateDict() → Map<String, float[]>¶

返回参数名到数值的映射

注意

数组是构造时的克隆，改数组不会回写文件，但会改变本实例持有的取值。

返回

不可修改的 Map，键为模块 namedParameters 的名称，值为 float 数组

final class ComputationGraph extends Module<Layout, Layout>¶

支持分支与残差连接的有向无环计算图

图由命名节点组成：输入节点用 addInputs 声明，MODULE 节点引用一个模块并只取单输入，ADD、SUB、MUL 节点只做张量运算、不持有参数。build 时按输入名与已登记节点名统计入度做 Kahn 拓扑排序，得到 executionOrder；排序覆盖不了全部节点就说明存在环，或存在无法从输入到达的节点，抛 IllegalArgumentException。

forward 按 executionOrder 逐节点求值，把输出存进以节点名为键的激活表：MODULE 取 inputs 的第一个名字调用 module.forward，ADD 与 MUL 从第二个输入起逐个累加或累乘，SUB 只接受两个输入做左减右。残差连接就是把同一个上游名字同时登记为模块节点的输入和 ADD 节点的一个输入，两条分支在 ADD 处合流。

多输入用 forwardMany(Tensor...) 或 forward(Map<String, Tensor>) 喂入，多输出分别以按 outputNames 顺序排列的数组和 LinkedHashMap 取回；单输入单输出的图可直接用 forward(Tensor)。求值全部基于普通张量算子，反向传播、train()/eval() 传播与 stateDict 存取都沿用模块体系。

示例

try (Basalt jsl = Basalt.create()) {
    ComputationGraph graph = ComputationGraph.builder()
            .addInputs("x")
            .addLayer("conv", new Conv2d(basalt.context(), 3, 8, 3, 3, 1, 1, 1, 1, true, 0L), "x")
            .addLayer("block", new Conv2d(basalt.context(), 8, 8, 3, 3, 1, 1, 1, 1, true, 0L), "conv")
            .addAddVertex("res", "conv", "block")   // 残差：把 conv 的输出与 block 的输出相加
            .setOutputs("res")
            .build();
    Tensor y = graph.forward(basalt.randn(2, 3, 16, 16));
    graph.close();
}

注意

ADD 与 MUL 至少需要两个输入，SUB 固定两个；名字与输入名或已有节点重名时在构造阶段抛 IllegalArgumentException。

注意

parameters() 按执行顺序拼接各模块的参数，ADD/SUB/MUL 节点不贡献参数；namedParameters 的键为“节点名.参数名”，因此 stateDict/loadStateDict 可直接用于整图存取。

注意

train()/eval() 与 close() 会传播到所有持有模块的节点；图本身不分配张量，中间结果由各算子按需创建。

注意

执行顺序在 build 时确定且不可变；vertices() 保持登记顺序，与 executionOrder() 的顺序不一定相同。

参见：basalt.torch.Module, basalt.torch.Sequential, basalt.torch.ComputationGraph.Builder

方法

builder() → Builderstatic¶

创建图构造器

返回

新的 Builder，输入、顶点与输出列表均为空

forward(Tensor input) → Tensor¶

单输入单输出的求值入口

注意

输出多于一个时只返回 outputNames 的第一个，其余输出仍会被计算但不返回。

参数

• input (Tensor) – 喂给唯一输入节点的张量，形状由首个模块决定

返回

第一个输出节点的张量

抛出

• IllegalStateException – 图的输入节点数不为 1

forwardMany(Tensor... inputs) → Tensor[]¶

按声明顺序喂入多个张量

注意

输入按声明顺序绑定，与参数名无关；多输出时数组长度等于输出节点数。

参数

• inputs (Tensor...) – 张量数组，长度与顺序必须与 addInputs 声明的输入名一致

返回

按 outputNames 顺序排列的输出张量数组

抛出

• IllegalArgumentException – 传入的张量个数与输入节点数不符

forward(Map<String, Tensor> inputs) → Map<String, Tensor>¶

以名字映射喂入并取回全部输出

注意

映射中不属于输入节点的多余键被忽略。

注意

这是另外两个 forward 重载的公共实现，单输入与多输入都经由它求值。

参数

• inputs (Map<String, Tensor>) – 输入名到张量的映射，必须覆盖全部输入名

返回

以输出名为键的 LinkedHashMap，顺序与 outputNames 一致

抛出

• BasaltParameterException – inputs 为 null

• IllegalArgumentException – 缺少某个必需的输入名

parameters() → List<Tensor>¶

按执行顺序汇总各模块参数

返回

所有 MODULE 节点参数的拼接列表；无参数节点不贡献元素

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回带节点名前缀的命名参数

返回

键为“节点名.参数名”的条目列表，顺序与执行顺序一致

train()¶

切换训练模式并向下传播

注意

先置自身的 training 标志，再对每个持有模块的节点调用 train()。

eval()¶

切换推理模式并向下传播

注意

先置自身的 training 标志，再对每个持有模块的节点调用 eval()，Dropout、BatchNorm2d 因此改走推理路径。

close()¶

关闭所有持有模块

注意

逐个调用子模块的 close()；外部传入的输入张量不归图所有，不会被释放。

inputNames() → List<String>¶

返回输入节点名

返回

不可变列表，顺序与 addInputs 的声明顺序一致

outputNames() → List<String>¶

返回输出节点名

返回

不可变列表，顺序与 setOutputs 的声明顺序一致

executionOrder() → List<String>¶

返回拓扑排序后的执行顺序

注意

顺序在 build 时固定，之后不再变化；未列入其中的节点不会被执行。

返回

节点名列表，每个节点都排在其全部输入之后

vertices() → Map<String, Vertex>¶

返回节点表

返回

节点名到 Vertex 的不可修改映射，保持登记顺序

final class Conv2d extends Module<BCHW, BCHW>¶

NCHW 上的二维卷积，支持分组、深度与膨胀

权重形状为 [outputChannels, inputChannels/groups, kernelHeight, kernelWidth]，即按分组切分输入通道的 OIHW；按 Glorot 均匀分布初始化，范围 ±sqrt(6 / (fanIn + fanOut))，其中 fanIn = (inputChannels/groups)·kH·kW、fanOut = (outputChannels/groups)·kH·kW，实际种子为 seed + 31·inputChannels + outputChannels。

偏置形状为 [1, outputChannels, 1, 1]，初始化为 0，在卷积结果上按通道广播相加；useBias 为 false 时不创建偏置，parameters() 与 namedParameters() 相应少一项。

输出空间尺寸为 (H + 2·paddingHeight − (dilationHeight·(kernelHeight − 1) + 1)) / strideHeight + 1，宽度同理，除法向下取整，结果任一维非正时拒绝执行。groups 只改变通道的分组方式，不改变尺寸公式；dilation 通过放大有效核尺寸参与计算，因此同样缩小输出。groups 等于 inputChannels 时为深度卷积，等于 1 时为普通卷积，介于两者之间为分组卷积。

示例

try (Basalt jsl = Basalt.create()) {
    Conv2d conv = new Conv2d(basalt.context(), 3, 16, 3, 3, 1, 1, 1, 1, true, 42L);
    Tensor input = basalt.randn(4, 3, 32, 32);
    Tensor output = conv.forward(input);
    // output 形状为 [4, 16, 32, 32]
    conv.close();
}

注意

构造器校验通道数、卷积核尺寸、步长、膨胀与 groups 均为正，且 inputChannels、outputChannels 能被 groups 整除；padding 不做校验，传入负值会缩小输出甚至导致输出尺寸非正。

注意

输入通道数必须能被 groups 整除，且 inputChannels/groups 必须等于权重的第二维，否则 Tensor.conv2d 抛 IllegalArgumentException。

注意

偏置通过广播相加实现，会额外产生一个计算图节点；不需要偏置时置 useBias=false 可省去该节点。

参见：basalt.torch.MaxPool2d, basalt.torch.Flatten, basalt.torch.Tensor#conv2d

构造器

Conv2d(NativeContext context, int inputChannels, int outputChannels, int kernelHeight, int kernelWidth)¶

默认 stride=1、padding=0、带偏置

注意

等价于 stride 1×1、padding 0×0、dilation 1×1、groups=1、useBias=true、seed=0 的完整构造器。

参数

• context (NativeContext) – 持有原生上下文，权重与偏置在其中分配

• inputChannels (int) – 输入通道数，必须为正

• outputChannels (int) – 输出通道数，即卷积核个数，必须为正

• kernelHeight (int) – 卷积核高度，必须为正

• kernelWidth (int) – 卷积核宽度，必须为正

抛出

• IllegalArgumentException – 任一通道数或卷积核尺寸小于等于 0

Conv2d(NativeContext context, int inputChannels, int outputChannels, int kernelHeight, int kernelWidth, int strideHeight, int strideWidth, int paddingHeight, int paddingWidth, boolean useBias, long seed)¶

完整指定步长、填充、偏置与种子

注意

等价于 dilation 1×1、groups=1 的完整构造器，即普通卷积。

注意

权重由 Tensor.random 以 ±sqrt(6/(fanIn+fanOut)) 的均匀分布生成，requiresGrad 为 true；偏置初值为 0。

注意

padding 大于核尺寸时输出可能变大，不会被拒绝。

参数

• context (NativeContext) – 持有原生上下文，权重与偏置在其中分配

• inputChannels (int) – 输入通道数，必须为正

• outputChannels (int) – 输出通道数，必须为正

• kernelHeight (int) – 卷积核高度，必须为正

• kernelWidth (int) – 卷积核宽度，必须为正

• strideHeight (int) – 垂直步长，必须为正；1 表示逐行滑动

• strideWidth (int) – 水平步长，必须为正

• paddingHeight (int) – 上下各补的零行数，参与输出尺寸计算；不校验，可为负

• paddingWidth (int) – 左右各补的零列数；不校验，可为负

• useBias (boolean) – 是否创建 [1, outputChannels, 1, 1] 的偏置张量

• seed (long) – 初始化种子基数，实际种子还会混入 inputChannels 与 outputChannels

抛出

• IllegalArgumentException – 任一通道数、卷积核尺寸或步长小于等于 0

Conv2d(NativeContext context, int inputChannels, int outputChannels, int kernelHeight, int kernelWidth, int strideHeight, int strideWidth, int paddingHeight, int paddingWidth, int dilationHeight, int dilationWidth, int groups, boolean useBias, long seed)¶

完整指定膨胀、分组、偏置与种子

注意

权重形状为 [outputChannels, inputChannels/groups, kernelHeight, kernelWidth]，按 ±sqrt(6/(fanIn+fanOut)) 的均匀分布初始化，fanIn 与 fanOut 均按每组通道数计算。

注意

实际种子为 seed + 31·inputChannels + outputChannels，与 groups、dilation 无关。

注意

groups 与 dilation 同时生效时，原生内核走逐元素累加路径，不再使用 im2col 加矩阵乘的快速路径。

参数

• context (NativeContext) – 持有原生上下文，权重与偏置在其中分配

• inputChannels (int) – 输入通道数，必须为正且能被 groups 整除

• outputChannels (int) – 输出通道数，必须为正且能被 groups 整除

• kernelHeight (int) – 卷积核高度，必须为正

• kernelWidth (int) – 卷积核宽度，必须为正

• strideHeight (int) – 垂直步长，必须为正

• strideWidth (int) – 水平步长，必须为正

• paddingHeight (int) – 上下各补的零行数；不校验，可为负

• paddingWidth (int) – 左右各补的零列数；不校验，可为负

• dilationHeight (int) – 垂直膨胀系数，必须为正；1 表示核元素相邻，2 表示间隔一个像素

• dilationWidth (int) – 水平膨胀系数，必须为正

• groups (int) – 分组数，必须为正；1 为普通卷积，等于 inputChannels 时为深度卷积

• useBias (boolean) – 是否创建 [1, outputChannels, 1, 1] 的偏置张量

• seed (long) – 初始化种子基数

抛出

• IllegalArgumentException – 任一通道数、卷积核尺寸、步长、膨胀或 groups 非正；或通道数不能被 groups 整除

方法

forward(Tensor input) → Tensor¶

对 NCHW 输入做卷积并叠加偏置

注意

偏置按通道广播，输出形状不因偏置改变。

注意

训练与推理行为相同，切换 train()/eval() 不影响结果。

参数

• input (Tensor) – 形状 [N, C, H, W]，C 必须等于 inputChannels

返回

[N, outputChannels, Hout, Wout] 的新张量

抛出

• IllegalArgumentException – 输入非四维、输入通道与权重不匹配，或输出空间尺寸非正

parameters() → List<Tensor>¶

返回权重与可选的偏置

返回

weight 在前；存在偏置时 bias 在后，否则长度为 1

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回 weight 与 bias 命名状态

返回

名称依次为 weight、bias；无偏置时只含 weight

weight() → Tensor¶

返回卷积核张量

返回

[outputChannels, inputChannels/groups, kernelHeight, kernelWidth] 的可训练张量

bias() → Tensor¶

返回偏置张量

注意

调用方需自行判空；无偏置时 parameters() 也不含该项。

返回

[1, outputChannels, 1, 1] 的可训练张量；useBias 为 false 时返回 null

inputChannels() → int¶

返回输入通道数

返回

构造时的 inputChannels，等于每组输入通道数乘以 groups

outputChannels() → int¶

返回输出通道数

返回

构造时的 outputChannels，即卷积核个数

kernelHeight() → int¶

返回卷积核高度

返回

构造时的 kernelHeight

kernelWidth() → int¶

返回卷积核宽度

返回

构造时的 kernelWidth

strideHeight() → int¶

返回垂直步长

返回

构造时的 strideHeight

strideWidth() → int¶

返回水平步长

返回

构造时的 strideWidth

paddingHeight() → int¶

返回垂直零填充行数

返回

构造时的 paddingHeight，可能为负

paddingWidth() → int¶

返回水平零填充列数

返回

构造时的 paddingWidth，可能为负

dilationHeight() → int¶

返回垂直膨胀系数

返回

构造时的 dilationHeight；1 表示标准卷积

dilationWidth() → int¶

返回水平膨胀系数

返回

构造时的 dilationWidth；1 表示标准卷积

groups() → int¶

返回分组数

返回

构造时的 groups；1 为普通卷积，等于 inputChannels 时为深度卷积

final class CosineAnnealingLR implements LRScheduler¶

余弦退火学习率调度器

按 lr_t = etaMin + 0.5·(baseLr − etaMin)·(1 + cos(π·t/T_max)) 调整优化器学习率，与 torch.optim.lr_scheduler.CosineAnnealingLR 一致；baseLr 在构造时从 optimizer.learningRate() 读取一次，之后不再跟随外部改动。

step() 先把内部计数器加一，再以 min(t, T_max) 代入公式，把结果写回优化器并返回；t 达到 T_max 后学习率停在 etaMin，继续调用只推进计数器、不再改变取值。

实现不含热重启：没有周期重置，也没有 restart 入口。需要周期性回升时，按周期新建调度器，或在重置点先 setLearningRate 回到 baseLr 再重新计时。

示例

try (Basalt jsl = Basalt.create()) {
    Linear model = new Linear(basalt.context(), 16, 4);
    AdamW optimizer = new AdamW(model.parameters(), 1e-3f);
    CosineAnnealingLR scheduler = new CosineAnnealingLR(optimizer, 100, 1e-5f);
    for (int epoch = 1; epoch <= 100; epoch++) {
        // ... 前向、反向
        optimizer.step();
        float lr = scheduler.step();
    }
}

注意

tMax 必须为正，etaMin 必须有限且非负，否则构造器抛 IllegalArgumentException；optimizer 为 null 抛 NullPointerException。

注意

currentLearningRate() 返回优化器当前的学习率，而不是按公式重算的值；外部改过优化器学习率后两者会不一致。

注意

计数器从 0 开始，第一次 step() 后 t = 1，因此要经过 T_max 次调用才降到 etaMin。

参见：basalt.torch.LRScheduler, basalt.torch.StepLR, basalt.torch.Optimizer#setLearningRate(float)

构造器

CosineAnnealingLR(Optimizer optimizer, int tMax)¶

etaMin 取 0 的余弦退火

注意

等价于 etaMin=0 的三参构造器，学习率最终退到 0。

注意

optimizer 为 null 时抛 NullPointerException。

参数

• optimizer (Optimizer) – 被调度的优化器，构造时读取其学习率作为 baseLr

• tMax (int) – 半个余弦周期的步数，必须为正

抛出

• IllegalArgumentException – tMax 小于等于 0

CosineAnnealingLR(Optimizer optimizer, int tMax, float etaMin)¶

指定周期长度与最小学习率

注意

etaMin 大于 baseLr 时曲线整体上移，学习率会随步数上升，构造器不做额外限制。

参数

• optimizer (Optimizer) – 被调度的优化器，构造时读取其学习率作为 baseLr

• tMax (int) – 半个余弦周期的步数，必须为正

• etaMin (float) – 学习率下界，必须有限且非负；0 表示可以退到 0

抛出

• IllegalArgumentException – tMax 小于等于 0，或 etaMin 为负、非有限

方法

step() → float¶

推进一步并写回学习率

注意

计数器先自增再钳到 tMax；超过 tMax 后返回值恒为 etaMin。

注意

每调用一次都会改写优化器的学习率，与是否执行过 optimizer.step() 无关。

返回

本步写入优化器的学习率

currentLearningRate() → float¶

返回优化器当前学习率

返回

optimizer.learningRate() 的当前值；尚未调用 step() 时等于构造时的 baseLr

currentStep() → int¶

返回已推进的步数

返回

从 0 开始、每次 step() 加一的计数，可超过 tMax

final class CrossEntropyLoss implements Loss¶

融合 softmax 与负对数似然的分类损失

loss(prediction, target) 直接转调 Tensor.crossEntropy：prediction 是 [N, C] 的未归一化 logits，target 是长度 N 的类索引（以 float 存储），返回形状 [1] 的标量，按 N 取平均。

前向用 log-sum-exp 稳定计算，反向一次得到 (softmax - onehot) / N，不需要显式构造 softmax 再取对数，省去中间张量与溢出风险。

注意

target 越界或非有限由原生层拒绝，抛 IllegalStateException；形状不符在 Java 侧抛 IllegalArgumentException。

注意

不支持类别权重、ignore_index 与 label smoothing，类别不平衡需自行加权或重采样。

注意

无状态、无可训练参数，可复用同一实例。

参见：basalt.torch.Loss, basalt.torch.MSELoss, basalt.torch.Tensor#crossEntropy(Tensor)

方法

loss(Tensor prediction, Tensor target) → Tensor¶

计算批量平均的交叉熵

注意

目标以 float 存储并被截断取整；越界由原生层抛 IllegalStateException。

参数

• prediction (Tensor) – 形状 [N, C] 的 logits，不做 softmax 预处理

• target (Tensor) – 元素个数为 N 的类索引张量，取值须落在 [0, C)

返回

形状 [1] 的标量损失张量，可继续调用 backward()

抛出

• IllegalArgumentException – logits 秩不为 2，或目标元素数不等于 N

final class DepthwiseSeparableConv2d extends Module<BCHW, BCHW>¶

深度可分离二维卷积，先逐通道再 1×1 混合

内部由两个 Conv2d 组成：depthwise 以 groups = inputChannels 做逐通道空间卷积，权重形状 [inputChannels, 1, kernelHeight, kernelWidth]，不设偏置；pointwise 是 1×1、步长 1、无填充的普通卷积，权重形状 [outputChannels, inputChannels, 1, 1]，是否带偏置由 useBias 决定。

forward 串联两段：先按 depthwise 的核、步长、填充、膨胀算出空间尺寸 Hout = (H + 2·pH − (dH·(kH − 1) + 1)) / sH + 1（向下取整），pointwise 不改变空间尺寸，最终输出 [N, outputChannels, Hout, Wout]。

参数按 depthwise 在前、pointwise 在后的顺序拼接；namedParameters 的键带 depthwise. 与 pointwise. 前缀，stateDict 因此能区分两段权重。

示例

try (Basalt jsl = Basalt.create()) {
    DepthwiseSeparableConv2d block =
            new DepthwiseSeparableConv2d(basalt.context(), 32, 64, 3, 2, 1, 42L);
    Tensor y = block.forward(basalt.randn(1, 32, 56, 56));
    // y 形状为 [1, 64, 28, 28]
    block.close();
}

注意

两段的种子分别是 seed + 17 与 seed + 53，再各自混入通道数，同一个 seed 也能得到不同的初值。

注意

参数校验完全交给两个 Conv2d：通道数、核尺寸、步长、膨胀非正时抛 IllegalArgumentException；inputChannels 与 outputChannels 无需相等。

注意

train()/eval() 与 close() 同时作用于两段；模块自身不持有额外张量。

注意

计算量约为普通 k×k 卷积的 1/outputChannels 加 1/(k×k)，代价是通道间信息只在 pointwise 段交换。

参见：basalt.torch.Conv2d, basalt.torch.BatchNorm2d, basalt.torch.Activation

构造器

DepthwiseSeparableConv2d(NativeContext context, int inputChannels, int outputChannels, int kernelSize, int stride, int padding, long seed)¶

方形核、对称步长与填充的构造器

注意

等价于 dilation 1×1、useBias=true 的完整构造器。

参数

• context (NativeContext) – 持有原生上下文，两段权重都在其中分配

• inputChannels (int) – 输入通道数，同时作为 depthwise 的分组数，必须为正

• outputChannels (int) – pointwise 的输出通道数，必须为正

• kernelSize (int) – depthwise 的方形核边长，必须为正

• stride (int) – 高宽共用的 depthwise 步长，必须为正

• padding (int) – 高宽共用的 depthwise 零填充宽度；不校验，可为负

• seed (long) – 种子基数，两段各自偏移后使用

抛出

• IllegalArgumentException – 通道数、核尺寸或步长非正

DepthwiseSeparableConv2d(NativeContext context, int inputChannels, int outputChannels, int kernelHeight, int kernelWidth, int strideHeight, int strideWidth, int paddingHeight, int paddingWidth, int dilationHeight, int dilationWidth, boolean useBias, long seed)¶

完整指定核、步长、填充、膨胀与偏置

注意

膨胀只作用于 depthwise；pointwise 固定 1×1、步长 1、无填充、groups=1。

注意

两段权重分别是 [inputChannels, 1, kH, kW] 与 [outputChannels, inputChannels, 1, 1]。

参数

• context (NativeContext) – 持有原生上下文，两段权重都在其中分配

• inputChannels (int) – 输入通道数，等于 depthwise 的分组数，必须为正

• outputChannels (int) – pointwise 的输出通道数，必须为正

• kernelHeight (int) – depthwise 核高度，必须为正

• kernelWidth (int) – depthwise 核宽度，必须为正

• strideHeight (int) – depthwise 垂直步长，必须为正

• strideWidth (int) – depthwise 水平步长，必须为正

• paddingHeight (int) – depthwise 垂直零填充；不校验，可为负

• paddingWidth (int) – depthwise 水平零填充；不校验，可为负

• dilationHeight (int) – depthwise 垂直膨胀系数，必须为正

• dilationWidth (int) – depthwise 水平膨胀系数，必须为正

• useBias (boolean) – 是否给 pointwise 段创建偏置；depthwise 段恒无偏置

• seed (long) – 种子基数，depthwise 用 seed + 17、pointwise 用 seed + 53

抛出

• IllegalArgumentException – 通道数、核尺寸、步长或膨胀非正

方法

forward(Tensor input) → Tensor¶

先逐通道卷积再 1×1 混合

注意

两段串联产生两个计算图节点，梯度沿链式法则先回传到 pointwise 输入，再进入 depthwise。

参数

• input (Tensor) – 形状 [N, inputChannels, H, W] 的张量

返回

[N, outputChannels, Hout, Wout] 的新张量，Hout、Wout 由 depthwise 的核、步长、填充与膨胀决定

抛出

• IllegalArgumentException – 输入非四维、通道数与权重不匹配，或输出空间尺寸非正

parameters() → List<Tensor>¶

返回两段参数

返回

depthwise 权重在前，随后是 pointwise 权重与可选偏置；无偏置时长度为 2

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回带段名前缀的命名参数

返回

depthwise.weight、pointwise.weight，带偏置时追加 pointwise.bias

depthwise() → Conv2d¶

返回逐通道卷积子模块

返回

groups 等于 inputChannels、useBias 为 false 的 Conv2d

pointwise() → Conv2d¶

返回 1×1 混合子模块

返回

核 1×1、groups 为 1 的 Conv2d，可能带偏置

train()¶

切换训练模式并传给两段

注意

先置自身的 training 标志，再对两段调用 train()；Conv2d 无状态，前向结果不受模式影响。

eval()¶

切换推理模式并传给两段

注意

先置自身的 training 标志，再对两段调用 eval()；对 Conv2d 而言与 train() 等价。

close()¶

关闭两段卷积

注意

释放两段的权重与偏置张量，之后模块不可再用。

final class Dropout<L extends Layout> extends Module<L, L>¶

倒置 dropout，推理期为恒等映射

训练模式下按概率 probability 把元素置零，并把保留的元素乘以 1 / (1 - probability)，使前向期望保持不变；probability 为 0 时直接返回输入的副本。

推理模式或 probability 为 0 时掩码全为 1，输出等于输入，因此同一模块可以在训练与推理之间反复切换而不改变调用方式。

掩码由原生内核用 xorshift64 从 seed 生成，给定种子完全确定；模块在每次 forward 之后把内部 seed 加一，所以连续调用得到不同掩码，而整条序列可复现。

注意

构造器校验 probability 落在 [0, 1)，越界抛 IllegalArgumentException；单参构造器的默认种子为 0x9E3779B97F4A7C15。

注意

掩码按元素个数生成，与形状无关；同一次前向与对应反向使用同一掩码。

注意

内部 seed 可变，同一实例不宜跨线程并发调用 forward。

参见：basalt.torch.Module#train(), basalt.torch.Tensor#dropout

构造器

Dropout(float probability)¶

以默认种子构造 dropout

参数

• probability (float) – 置零概率，取值 [0, 1)；0 表示恒等

抛出

• IllegalArgumentException – probability 不在 [0, 1) 内，含 NaN

Dropout(float probability, long seed)¶

指定掩码种子构造 dropout

注意

种子为 0 时原生侧会先与 1 按位或，实际状态序列与传 1 相同。

参数

• probability (float) – 置零概率，取值 [0, 1)

• seed (long) – 掩码生成种子，每次 forward 后自增

抛出

• IllegalArgumentException – probability 不在 [0, 1) 内，含 NaN

方法

forward(Tensor input) → Tensor¶

训练时随机置零并缩放，推理时恒等

注意

每次调用都会推进内部种子，因此同一输入连续两次前向的结果不同。

参数

• input (Tensor) – 任意形状的张量，掩码按元素个数生成

返回

与输入同形状的新张量

parameters() → List<Tensor>¶

返回空参数列表

返回

恒为空的不可变列表

final class Embedding extends Module<BC, BSF>¶

词表查找表，索引以 f32 存整数

权重形状为 [vocabularySize, dimension]，按 ±1/sqrt(dimension) 的均匀分布初始化，实际种子为 seed + 41·vocabularySize + dimension。

forward 把输入中的每个元素当作行号，取出对应的 dimension 维向量，输出形状为输入形状末尾追加一维 dimension；输入可以是任意秩，不限于 [batch, sequence]。

索引以 f32 存储，原生内核校验其有限且落在 [0, vocabularySize) 内，随后截断为整数行号再拷贝整行权重。

注意

索引应为整数值的浮点数；带小数部分时会被截断取整，不报错。

注意

vocabularySize 与 dimension 必须为正，否则张量分配阶段抛 IllegalArgumentException。

注意

反向传播按索引做散射累加，同一索引在输入中出现多次时梯度相加。

参见：basalt.torch.Linear, basalt.torch.Tensor#embedding

构造器

Embedding(NativeContext context, int vocabularySize, int dimension)¶

以 seed=0 构造查找表

注意

等价于四参构造器传入 seed=0。

参数

• context (NativeContext) – 持有原生上下文，权重在其中分配

• vocabularySize (int) – 词表大小，即权重行数，必须为正

• dimension (int) – 每个索引对应的向量维度，必须为正

Embedding(NativeContext context, int vocabularySize, int dimension, long seed)¶

指定初始化种子构造查找表

注意

初始化幅度随 dimension 增大而收缩，dimension=1 时为 ±1。

参数

• context (NativeContext) – 持有原生上下文，权重在其中分配

• vocabularySize (int) – 词表大小，必须为正

• dimension (int) – 向量维度，必须为正

• seed (long) – 初始化种子基数，实际种子混入 vocabularySize 与 dimension

方法

forward(Tensor input) → Tensor¶

按索引查表，输出追加一维

注意

输出不复制权重之外的数据，行内容与权重对应行逐元素相同。

参数

• input (Tensor) – 元素为整数值的浮点张量，任意秩，每个元素是一个行号

返回

形状为 input.shape() 末尾追加 dimension 的新张量

抛出

• BasaltRuntimeException – 索引非有限，或落在 [0, vocabularySize) 之外（由原生内核校验）

parameters() → List<Tensor>¶

返回权重张量

返回

只含 weight 的单元素列表

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回名为 weight 的状态项

返回

只含 weight 的单元素列表

weight() → Tensor¶

返回查找表权重

返回

[vocabularySize, dimension] 的可训练张量

final class Flatten extends Module<BCHW, BC>¶

把 NCHW 展平为 [batch, C·H·W]

forward 读取输入形状，要求秩为 4，然后返回 reshape(shape[0], shape[1] * shape[2] * shape[3]) 的结果；通道、高度、宽度按行主序依次拼接，不做转置。

模块没有参数，类型上把 BCHW 布局转换为 BC 布局，因此在 Sequential 中只有展平之后才允许接 Linear 这类二维模块。

注意

输入秩不为 4 时抛 IllegalArgumentException，不接受已经展平的二维输入。

注意

reshape 会新分配张量并登记一个计算图节点，反向时按逆映射回传梯度，输出与输入不共享存储。

参见：basalt.torch.Sequential, basalt.torch.Linear, basalt.torch.Tensor#reshape

方法

forward(Tensor input) → Tensor¶

把四维输入展平为二维

参数

• input (Tensor) – 形状 [N, C, H, W] 的张量

返回

[N, C·H·W] 的新张量

抛出

• IllegalArgumentException – 输入秩不为 4

parameters() → List<Tensor>¶

返回空参数列表

返回

恒为空的不可变列表

final class Graph implements AutoCloseable¶

原生自动求导图的 Java 侧句柄：反向、清图与统计

Tensor.backward() 是最常用的反向入口；Graph 面向需要显式控制图生命周期的场景：在独立前向之间清图、查看已记录的节点数、读取上下文级存活统计。构造器与 of 等价，只校验 context 非空。

close() 仅等价于 clear()，不关闭底层 NativeContext；调用后上下文仍可继续创建张量、执行算子并重新累积节点。

注意

图是追加式的：每次指令队列执行都会登记一个节点，与输入是否 requiresGrad 无关；训练循环必须每步 clear()，否则节点与中间张量持续累积。

注意

backward 不清图也不清零梯度；清图之后原根张量不再是当前图的产物，再调用 backward 会抛 IllegalStateException。

注意

关闭图中任一被引用的张量会隐式清空整张图，这是原生 release_tensor 的既定语义。

注意

size() 与 stats() 都是即时查询，各触发一次 JNI 调用，不宜在热循环中频繁调用。

参见：basalt.torch.Tensor#backward(), basalt.torch.Optimizer, basalt.nativebridge.NativeContext

构造器

Graph(NativeContext context)¶

在给定上下文上创建图句柄

注意

创建动作本身不改动图内容；多个 Graph 实例操作的是同一张图。

参数

• context (NativeContext) – 原生上下文，图句柄只是它的一个视图，不持有额外原生资源

抛出

• BasaltRuntimeException – context 为 null

方法

of(NativeContext context) → Graphstatic¶

构造器的静态工厂形式

参数

• context (NativeContext) – 原生上下文

返回

绑定该上下文的新图句柄

抛出

• BasaltRuntimeException – context 为 null

backward(Tensor root)¶

从根张量反向传播并累积梯度

初始余切为全 1 向量，沿拓扑序逆序访问可达节点，把梯度累加到 requiresGrad 为 true 的张量上。

注意

根张量未被当前图产生时，原生层抛 IllegalStateException。

注意

梯度以累加方式写入，多次调用会叠加而非覆盖。

参数

• root (Tensor) – 反向起点，必须是当前计算图的产物；通常是形状 [1] 的损失

抛出

• BasaltRuntimeException – root 为 null

• IllegalArgumentException – root 属于另一个上下文

clear()¶

丢弃全部已记录的节点

清空节点表与产出、消费索引，之后新建的张量可以重新作为图节点输出。

注意

只清图，不释放张量，也不清梯度；正在被引用的中间张量仍然有效。

size() → int¶

返回当前图中的节点数

注意

节点数与算子调用次数一致，可用于确认清图是否生效。

返回

自上次 clear() 以来登记的节点个数

stats() → long[]¶

返回上下文级的存活统计

注意

字节数是原生侧实际占用，可用于定位张量泄漏。

返回

长度为 3 的数组，依次为存活张量数、存活字节数、图节点数

close()¶

清空计算图，等价于 clear()

注意

不关闭上下文；配合 try-with-resources 可让一步训练结束即清图。

final class GroupNorm<L extends Layout> extends Module<L, L>¶

按通道分组归一化，结果与批量大小无关

构造时按 numChannels 分配两个一维可训练张量 gamma 与 beta，形状均为 [numChannels]，分别填充 1 与 0；默认 epsilon 为 1e-5f。类本身只保存配置，不持有滑动统计量。

forward 转调 Tensor.groupNorm：把第二维的 C 个通道均分为 numGroups 组，每组包含 numChannels / numGroups 个通道及其全部空间位置；在该组内计算均值与方差，除以组内元素数，归一化后按通道应用 gamma、beta。

其余维度按行主序合并为空间维，统计既不跨样本也不跨分组；输出与输入同形状，因此可以直接接在 Conv2d 之后。与 BatchNorm2d 不同，训练与推理行为完全一致。

注意

numGroups 与 numChannels 都必须为正，且 numChannels 必须能被 numGroups 整除，否则构造器抛 IllegalArgumentException。

注意

gamma 与 beta 的形状是 [numChannels]，广播按通道进行；二者元素数与输入第二维不一致时，forward 抛 IllegalArgumentException。

注意

numGroups 等于 1 时退化为按样本的层归一化，等于 numChannels 时退化为逐通道归一化。

注意

输入秩为 2 时空间维按 1 处理，只对通道向量分组统计。

注意

epsilon 直接参与 sqrt(var + epsilon)，构造阶段不校验取值。

参见：basalt.torch.LayerNorm, basalt.torch.BatchNorm2d, basalt.torch.Tensor#groupNorm

构造器

GroupNorm(NativeContext context, int numGroups, int numChannels)¶

以 epsilon=1e-5 构造分组归一化

注意

等价于四参构造器传入 1e-5f。

注意

gamma、beta 由 Tensor.parameter 创建，requiresGrad 为 true，随后分别填充 1 与 0。

参数

• context (NativeContext) – 持有原生上下文，gamma 与 beta 在其中分配

• numGroups (int) – 分组数，必须为正且能整除 numChannels

• numChannels (int) – 通道数，同时是 gamma、beta 的长度

抛出

• IllegalArgumentException – numGroups 或 numChannels 非正，或 numChannels 不能被 numGroups 整除

GroupNorm(NativeContext context, int numGroups, int numChannels, float epsilon)¶

显式指定数值稳定项 epsilon

注意

epsilon 为负且组内方差接近 0 时结果可能非有限，构造阶段不会拦截。

注意

gamma 与 beta 的形状是 [numChannels] 的一维张量，不是 [1, C, 1, 1]。

参数

• context (NativeContext) – 持有原生上下文，gamma 与 beta 在其中分配

• numGroups (int) – 分组数，必须为正且能整除 numChannels

• numChannels (int) – 通道数，同时是 gamma、beta 的长度

• epsilon (float) – 加在方差上的稳定项，避免除零；构造阶段不校验正负

抛出

• IllegalArgumentException – numGroups 或 numChannels 非正，或 numChannels 不能被 numGroups 整除

方法

forward(Tensor input) → Tensor¶

对输入按通道分组归一化

注意

统计范围是每个样本、每个分组内的全部通道与空间位置，不跨样本。

注意

原生层还要求 gamma、beta 的形状恰为 [C]，否则报形状不匹配。

参数

• input (Tensor) – 秩至少为 2 的张量，第二维必须等于 numChannels，其余维度视为空间维

返回

与输入同形状的新张量，已登记计算图节点

抛出

• IllegalArgumentException – 输入秩小于 2 或第二维不能被 numGroups 整除；gamma、beta 的元素数与第二维不一致

parameters() → List<Tensor>¶

返回 gamma 与 beta

返回

长度为 2 的列表，顺序为 gamma、beta

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回 gamma、beta 两个命名状态

返回

顺序固定为 gamma、beta

numGroups() → int¶

返回分组数

返回

构造时传入的 numGroups

numChannels() → int¶

返回通道数

返回

构造时传入的 numChannels，即 gamma、beta 的长度

gamma() → Tensor¶

返回缩放参数

返回

[numChannels] 的可训练张量，初值为全 1

beta() → Tensor¶

返回平移参数

返回

[numChannels] 的可训练张量，初值为全 0

final class L1Loss implements Loss¶

平均绝对误差损失（L1）

loss(prediction, target) 转调 Tensor.l1Loss：先计算 prediction − target，逐元素取绝对值后求算术平均，返回 [1] 标量，可直接调用 backward()。

归约方式是全部元素的算术平均，除以元素总数而非 batch 大小；实现上把绝对差 reshape 成 [1, n]，与长度为 n、元素为 1/n 的列向量做矩阵乘完成求和。

注意

形状不一致抛 IllegalArgumentException，消息为 L1Loss expects equal shapes。

注意

除差值恰为 0 的位置外，反向梯度为 sign(prediction − target)/n，对 prediction 与 target 两侧同时累积。

注意

与 MSELoss 一样不做广播；需要按样本平均时先自行 reshape。

参见：basalt.torch.Loss, basalt.torch.MSELoss, basalt.torch.Tensor#l1Loss

方法

loss(Tensor prediction, Tensor target) → Tensor¶

计算逐元素绝对误差的均值

注意

不做广播：即使形状可广播也会被拒绝。

注意

归约除以全部元素个数，而非 batch 大小。

参数

• prediction (Tensor) – 预测张量，任意秩

• target (Tensor) – 目标张量，形状必须与 prediction 完全相同

返回

形状 [1] 的标量损失张量

抛出

• IllegalArgumentException – 两个张量形状不完全相同

interface LRScheduler¶

学习率调度器的统一契约

实现类在内部维护步数计数：step() 推进一步并算出新学习率，通过 Optimizer.setLearningRate 写回优化器，同时返回该值；currentLearningRate() 与 currentStep() 提供只读查询。三个方法都不带参数。

调用时机由使用方决定，每完成一个 epoch 或每个优化步后调用一次 step()，实现类不区分钟与步。StepLR 按 currentStep / stepSize 计算衰减次数，CosineAnnealingLR 按 min(t, tMax) 计算余弦相位，两者都把一次 step() 计为一格。

构造时实现类通常读取 optimizer.learningRate() 作为基准学习率，此后只按自己的公式写回；换基准需要新建调度器。

注意

step() 直接改写优化器学习率，与是否调用过 optimizer.step() 无关。

注意

写回的学习率必须有限且非负，否则 Optimizer.setLearningRate 抛 IllegalArgumentException。

注意

实现类在构造时绑定单个 Optimizer，同一实例不应跨优化器复用。

注意

接口不声明异常，也不规定步数上限；计数只增不减，没有重置入口。

参见：basalt.torch.StepLR, basalt.torch.CosineAnnealingLR, basalt.torch.Optimizer#setLearningRate(float)

final class LayerNorm<L extends Layout> extends Module<L, L>¶

沿最后一维做层归一化

forward 在最后一维上计算均值与方差，加 epsilon 后开方做归一化，再逐元素乘以 gamma、加上 beta；其余维度各自视为独立样本，不参与统计。

gamma 与 beta 的形状取自 normalizedShape，分别初始化为 1 与 0。Tensor.layerNorm 只要求两者的元素个数等于输入最后一维，因此通常传单个维度值即可。

注意

构造器校验 normalizedShape 非空；维度必须为正，否则张量分配阶段抛 IllegalArgumentException。

注意

层归一化与批量无关，没有滑动统计量，训练与推理行为完全相同。

注意

gamma 与 beta 的元素个数和输入最后一维不一致时，forward 抛 IllegalArgumentException。

参见：basalt.torch.BatchNorm2d, basalt.torch.Tensor#layerNorm

构造器

LayerNorm(NativeContext context, int... normalizedShape)¶

以 epsilon=1e-5 构造层归一化

注意

等价于传入 1e-5f 的三参构造器。

参数

• context (NativeContext) – 持有原生上下文，gamma 与 beta 在其中分配

• normalizedShape (int...) – 被归一化的维度尺寸，通常只传最后一维长度；不得为空

抛出

• IllegalArgumentException – normalizedShape 为空

LayerNorm(NativeContext context, float epsilon, int... normalizedShape)¶

显式指定数值稳定项

注意

epsilon 为负且方差接近 0 时可能得到非有限值，构造阶段不会拦截。

参数

• context (NativeContext) – 持有原生上下文，gamma 与 beta 在其中分配

• epsilon (float) – 加在方差上的稳定项，避免除零；不校验正负

• normalizedShape (int...) – gamma 与 beta 的形状，元素个数需等于输入的最后一维

抛出

• IllegalArgumentException – normalizedShape 为空

方法

forward(Tensor input) → Tensor¶

沿最后一维归一化并仿射

参数

• input (Tensor) – 任意秩张量，最后一维长度需等于 gamma 的元素个数

返回

与输入同形状的新张量

抛出

• IllegalArgumentException – gamma 或 beta 的元素个数与输入最后一维不一致

parameters() → List<Tensor>¶

返回 gamma 与 beta

返回

长度为 2 的列表，顺序为 gamma、beta

epsilon() → float¶

返回数值稳定项

返回

构造时传入的 epsilon；双参构造器为 1e-5f

gamma() → Tensor¶

返回缩放参数

返回

形状为 normalizedShape 的可训练张量，初值为全 1

beta() → Tensor¶

返回平移参数

返回

形状为 normalizedShape 的可训练张量，初值为全 0

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回名为 gamma、beta 的状态项

返回

顺序固定为 gamma、beta

final class Linear extends Module<BC, BC>¶

全连接层 y = xW + b，只投影最后一维

权重形状为 [inputSize, outputSize]，按 ±sqrt(6 / (inputSize + outputSize)) 的均匀分布初始化，实际种子为 seed + 17·inputSize + outputSize；偏置形状为 [1, outputSize]，初始化为 0。

forward 只要求输入的最后一维等于 inputSize：秩为 2 时直接做二维矩阵乘；秩大于 2 或秩为 1 时先把前置维度合并成行，乘完再恢复原形状，仅把最后一维替换为 outputSize。

因此 [batch, time, channels] 这类输入会被逐位置投影，前置维度保持不变，与 PyTorch 的 Linear 语义一致。

注意

构造器不校验 inputSize、outputSize 的取值，非正尺寸在张量分配阶段抛 IllegalArgumentException。

注意

useBias 为 false 时 bias() 返回 null，parameters() 与 namedParameters() 只含 weight。

注意

形状检查发生在前向，最后一维不匹配时异常信息包含期望值与实际值。

参见：basalt.torch.LayerNorm, basalt.torch.Sequential, basalt.torch.Tensor#matMul

构造器

Linear(NativeContext context, int inputSize, int outputSize)¶

默认带偏置、seed=0 构造

注意

等价于五参构造器传入 useBias=true、seed=0。

参数

• context (NativeContext) – 持有原生上下文，权重与偏置在其中分配

• inputSize (int) – 输入特征数，即输入最后一维的长度，必须为正

• outputSize (int) – 输出特征数，必须为正

Linear(NativeContext context, int inputSize, int outputSize, boolean useBias, long seed)¶

指定偏置开关与初始化种子

注意

权重由 Tensor.random 以均匀分布生成，requiresGrad 为 true；偏置初值为 0。

参数

• context (NativeContext) – 持有原生上下文，权重与偏置在其中分配

• inputSize (int) – 输入特征数，必须为正

• outputSize (int) – 输出特征数，必须为正

• useBias (boolean) – 是否创建 [1, outputSize] 的偏置张量

• seed (long) – 初始化种子基数，实际种子混入 inputSize 与 outputSize

方法

forward(Tensor input) → Tensor¶

沿最后一维做仿射变换

注意

秩大于 2 时前置维度被合并成行，输出仍保留这些维度。

参数

• input (Tensor) – 最后一维需等于 inputSize 的张量，秩不限

返回

形状与输入相同，仅最后一维替换为 outputSize

抛出

• IllegalArgumentException – 输入的最后一维不等于 inputSize

parameters() → List<Tensor>¶

返回权重与可选的偏置

返回

weight 在前；存在偏置时 bias 在后，否则长度为 1

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回 weight 与 bias 命名状态

返回

名称依次为 weight、bias；无偏置时只含 weight

weight() → Tensor¶

返回权重张量

返回

[inputSize, outputSize] 的可训练张量

bias() → Tensor¶

返回偏置张量

返回

[1, outputSize] 的可训练张量；构造时 useBias 为 false 则返回 null

interface Loss¶

可微标量损失的统一接口

实现类接收预测张量与目标张量，返回形状 [1] 的标量张量，可直接调用 backward()。接口本身不约束形状、归约方式与目标编码，这些由具体实现或底层 Tensor 算子决定。

注意

实现通常无状态、无可训练参数，可作为单例复用。

注意

接口不声明异常，形状错误以 IllegalArgumentException 或 IllegalStateException 的形式在运行期暴露。

参见：basalt.torch.MSELoss, basalt.torch.CrossEntropyLoss, basalt.torch.Tensor#backward()

final class MSELoss implements Loss¶

均方误差损失

loss(prediction, target) 转调 Tensor.mse：要求两个张量形状完全相同，返回 [1] 标量。归约方式是全部元素的算术平均，除以元素总数而非 batch 大小。

注意

形状不一致抛 IllegalArgumentException，消息为 MSE expects equal shapes。

注意

反向梯度为 2*(prediction-target)/n，对 prediction 与 target 两侧同时累积。

参见：basalt.torch.Loss, basalt.torch.CrossEntropyLoss, basalt.torch.Tensor#mse(Tensor)

方法

loss(Tensor prediction, Tensor target) → Tensor¶

计算逐元素平方误差的均值

注意

不做广播：即使形状可广播也会被拒绝。

参数

• prediction (Tensor) – 预测张量，任意秩

• target (Tensor) – 目标张量，形状必须与 prediction 完全相同

返回

形状 [1] 的标量损失张量

抛出

• IllegalArgumentException – 两个张量形状不完全相同

final class MaxPool2d extends Module<BCHW, BCHW>¶

NCHW 上的最大池化，不支持填充

forward 调用 Tensor.maxPool2d，输出尺寸为 (H − kernelHeight) / strideHeight + 1，除法向下取整，宽度同理；窗口内取最大值，反向时梯度只回传到每个窗口最大值所在的位置。

模块没有 padding 参数，实现也不做隐式填充，边缘不足一个完整窗口的位置会被丢弃；单参构造器令 stride 等于 kernel，得到互不重叠的窗口。

注意

构造器不校验 kernel 与 stride 的取值；stride 传 0 会在原生内核中除零。

注意

输入秩不为 4，或任一输出维非正时抛 IllegalArgumentException。

注意

模块无参数，parameters() 恒为空列表。

参见：basalt.torch.Conv2d, basalt.torch.Tensor#maxPool2d

构造器

MaxPool2d(int kernelSize)¶

方形核且步长等于核尺寸

注意

等价于四参构造器传入四个相同的值。

参数

• kernelSize (int) – 窗口边长，同时用作垂直与水平步长

MaxPool2d(int kernelHeight, int kernelWidth, int strideHeight, int strideWidth)¶

分别指定核尺寸与步长

注意

stride 小于 kernel 时窗口重叠，大于 kernel 时跳过部分位置。

参数

• kernelHeight (int) – 窗口高度，必须为正才会得到有效输出

• kernelWidth (int) – 窗口宽度，必须为正才会得到有效输出

• strideHeight (int) – 垂直步长，不校验取值

• strideWidth (int) – 水平步长，不校验取值

方法

forward(Tensor input) → Tensor¶

对 NCHW 输入做最大池化

参数

• input (Tensor) – 形状 [N, C, H, W] 的张量

返回

[N, C, Hout, Wout] 的新张量

抛出

• IllegalArgumentException – 输入非四维，或任一输出维非正

parameters() → List<Tensor>¶

返回空参数列表

返回

恒为空的不可变列表

kernelHeight() → int¶

返回窗口高度

返回

构造时传入的 kernelHeight

kernelWidth() → int¶

返回窗口宽度

返回

构造时传入的 kernelWidth

strideHeight() → int¶

返回垂直步长

返回

构造时传入的 strideHeight

strideWidth() → int¶

返回水平步长

返回

构造时传入的 strideWidth

abstract class Module<I extends Layout, O extends Layout> implements AutoCloseable¶

带布局类型参数的模块基类，统一参数与训练态

两个类型参数 I、O 是幽灵布局标记，只在编译期参与 Sequential 的连线检查，运行期不做形状校验；forward 内部仍需按实际形状处理输入，形状不符时由具体模块抛出异常。

子类实现 forward 与 parameters；namedParameters 默认给每个参数起 param0…paramN 的位置名，stateDict 以此为键。train()/eval() 只翻转基类中的布尔标志，是否据此切换行为由子类 forward 自行读取。

注意

构造后默认处于训练态，即 training=true。

注意

apply 只做空值检查后转调 forward，本身不改变语义。

注意

stateDict 返回的是数据副本，修改返回值不影响模块；loadStateDict 缺少条目抛 IllegalArgumentException，多余条目被忽略。

注意

close() 关闭 parameters() 中的每个张量；容器模块应保证同一张量不被重复关闭，句柄关闭本身是幂等的。

注意

基类不持有张量，也不做参数注册：参数必须由子类保存在字段中并通过 parameters() 暴露。

参见：basalt.torch.Sequential, basalt.torch.Linear, basalt.torch.Optimizer

方法

forward(Tensor input) → Tensorabstract¶

子类实现的前向计算

注意

抽象方法，不做空值检查；需要防御性检查时用 apply。

参数

• input (Tensor) – 输入张量，形状约定由具体模块决定

返回

输出张量，通常是新分配的张量并已登记计算图节点

parameters() → List<Tensor>abstract¶

返回模块持有的可训练参数

注意

无参数模块返回空列表；返回的列表不应被调用方修改。

返回

参数列表，顺序稳定，供 Optimizer 与 zeroGrad 使用

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回带名字的参数条目

注意

默认实现只按位置命名；需要稳定语义名时子类应覆盖，如 Linear 使用 weight 与 bias。

返回

名字到张量的条目列表，默认命名为 param0、param1…

apply(Tensor input) → Tensor¶

校验输入非空后调用 forward

注意

模块组合时统一走 apply 可把空值错误定位在边界上。

参数

• input (Tensor) – 输入张量

返回

forward 的返回值

抛出

• BasaltRuntimeException – input 为 null

train()¶

把模块切到训练态

注意

只设置基类标志，不递归；容器模块会覆盖此方法以传播到子模块。

eval()¶

把模块切到推理态

注意

只设置基类标志，不递归；Dropout、BatchNorm2d 等模块在 forward 中读取该标志。

isTraining() → boolean¶

返回当前是否为训练态

返回

训练态为 true，推理态为 false；新模块默认为 true

zeroGrad()¶

清零全部参数的梯度

注意

遍历 parameters() 逐个调用 Tensor.zeroGrad，不清计算图。

stateDict() → Map<String, float[]>¶

把参数导出为名字到数据的映射

注意

副本可安全序列化或修改；键来自 namedParameters。

返回

保持插入顺序的映射，值是参数的完整数据副本

loadStateDict(Map<String, float[]> state)¶

按名字把数据写回参数

注意

多余条目被静默忽略；写入是原地覆盖，不影响计算图结构。

参数

• state (Map<String, float[]>) – 名字到数据的映射，必须覆盖当前模块的全部参数名

抛出

• IllegalArgumentException – 某个数组长度与对应参数元素数不一致

close()¶

关闭模块持有的全部参数张量

注意

逐个调用 Tensor.close；参数被当前计算图引用时会连带清空整张图。

final class MultiheadAttention extends Module<BSF, BSF>¶

多头自注意力，输入 [batch, time, channels]

构造时创建 query、key、value、output 四个 dimension 到 dimension 的 Linear，均带偏置，并创建作用于注意力权重的 Dropout；缩放系数是值为 1/sqrt(dimension / heads) 的标量张量。

forward 先把每个投影结果 reshape 成 [B, T, H, dh]：Q 与 V 经 permute(0, 2, 1, 3) 变为 [B·H, T, dh]，K 经 permute(0, 2, 3, 1) 变为 [B·H, dh, T]，再做批量矩阵乘得到 [B·H, T, T] 的分数。

分数乘缩放系数，causal 为真时把 column > row 的位置加上 -1e9，随后沿最后一维 softmax、过 dropout，再与 V 相乘得 [B·H, T, dh]，reshape 成 [B, H, T, dh] 并 permute(0, 2, 1, 3) 还原为 [B, T, C]，最后经输出投影。

示例

try (Basalt jsl = Basalt.create()) {
    MultiheadAttention attention =
            new MultiheadAttention(basalt.context(), 64, 8, true, 0.1f, 7L);
    Tensor input = basalt.randn(2, 16, 64);
    attention.train();
    Tensor output = attention.forward(input);
    // output 形状为 [2, 16, 64]
    attention.eval();
    attention.close();
}

注意

dimension 必须能被 heads 整除，否则构造与前向都抛 IllegalArgumentException；heads 本身不做正值校验。

注意

因果掩码只屏蔽未来位置（严格上三角），每次前向都按 time×time 新建一个掩码张量。

注意

train()/eval() 会同步切换内部 dropout；parameters() 只含四个投影的权重与偏置，不含缩放常量。

注意

namedParameters 的名称前缀为 query.、key.、value.、output.。

参见：basalt.torch.SelfAttention, basalt.torch.TransformerBlock, basalt.torch.Linear

构造器

MultiheadAttention(NativeContext context, int dimension, int heads, boolean causal, float dropoutProbability, long seed)¶

构造多头注意力与四个投影

注意

每个投影都是带偏置的方阵 Linear，因此参数总量为 4·(dimension² + dimension)。

参数

• context (NativeContext) – 持有原生上下文，四个投影、缩放张量与 dropout 都在其中

• dimension (int) – 模型维度，同时是输入最后一维与各投影的输入输出尺寸

• heads (int) – 注意力头数，dimension 必须能被其整除

• causal (boolean) – 是否加因果掩码，为 true 时位置 t 只能看到不晚于 t 的位置

• dropoutProbability (float) – 注意力权重上的置零概率，取值 [0, 1)

• seed (long) – 种子基数，四个投影分别偏移 101、211、307、401，dropout 偏移 503

抛出

• IllegalArgumentException – dimension 不能被 heads 整除

方法

forward(Tensor input) → Tensor¶

计算多头自注意力输出

注意

Q、K、V 都来自同一输入，属于自注意力，不提供分离的 key/value 输入。

注意

注意力权重在 softmax 之后、与 V 相乘之前经过 dropout。

参数

• input (Tensor) – 形状 [batch, time, channels]，channels 需能被 heads 整除

返回

[batch, time, channels]，最后一维与输入相同

抛出

• IllegalArgumentException – 输入不是三维，或通道数不能被 heads 整除

parameters() → List<Tensor>¶

返回四个投影的参数

返回

query、key、value、output 的权重与偏置依次拼接

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回带模块前缀的命名参数

返回

名称形如 query.weight、query.bias、key.weight 等

train()¶

切换到训练模式并同步 dropout

注意

置位自身标志后调用内部 dropout 的 train()，其余投影没有模式相关行为。

eval()¶

切换到推理模式并同步 dropout

注意

推理时 dropout 为恒等映射，注意力权重不再被置零。

close()¶

释放子模块与缩放张量

注意

四个投影、dropout 与缩放张量都由本模块持有并一并释放。

final class NativeMemoryEvent extends Event¶

原生上下文内存快照的 JFR 事件

BasaltMetrics.snapshot(NativeContext) 在返回 MetricsSnapshot 的同时填充并 commit 该事件，带出存活张量数、离堆字节数、计算图节点数与 JVM 堆已用量。

JFR 是 JDK 内置的事件框架，可用 jcmd <pid> JFR.start 在运行时打开采集，不必修改代码或重启进程；默认不记录，关闭时每次提交只付出一次 shouldCommit() 标志位读取。

事件粒度是调用方主动采样的快照，不是每个算子一条；算子级明细由 OperatorExecutionEvent 承担。

注意

事件名为 basalt.NativeMemory，标签 Basalt Native Memory，类别 {Basalt, Observability}。

注意

三个原生计数取自 NativeContext.stats()；上下文为 null、已关闭或统计抛异常时全部置 0，事件仍照常提交。

注意

heapUsedBytes 为 Runtime.totalMemory() - freeMemory()，含尚未回收的垃圾，因此高于真实存活对象。

注意

字段全是 long 原始类型，JFR 以定长记录写入线程本地缓冲；未启用采集时字段赋值与 commit 一并跳过。

注意

不启动采样线程，采样节奏完全由调用方决定。

参见：basalt.BasaltMetrics, basalt.torch.OperatorExecutionEvent, basalt.nativebridge.NativeContext

字段

liveTensors()¶

快照时刻存活的原生张量数

注意

取自 NativeContext.stats() 的第一个元素；上下文不可用时为 0。

liveBytes()¶

存活张量占用的离堆字节数

注意

NativeContext.stats() 的第二个元素，按张量数据缓冲累计，不含 JVM 对象头。

graphNodes()¶

当前计算图的节点数

注意

NativeContext.stats() 的第三个元素；未清图时随算子调用单调增长。

heapUsedBytes()¶

JVM 堆已用字节数

注意

Runtime.totalMemory() - freeMemory()，含未回收垃圾，不等于存活对象大小。

final class NativeQueueEvent extends Event¶

每次原生指令队列执行产生的 JFR 事件

CommandQueueBuilder.execute 在每次批量 JNI 调用结束时填充并 commit 该事件，记录本次队列的算子条数、线程提示、载荷字节数与墙钟耗时。

JFR 是 JDK 内置的事件框架，可在运行时用 jcmd <pid> JFR.start 打开采集，无需修改代码或重启进程，因此这条事件流在默认配置下始终可用。

注意

事件名为 basalt.NativeQueue，标签 Native Command Queue，类别 {Basalt, Native}。

注意

bytes 等于 commands * 128，对应每条指令 128 字节的定长线格式。

注意

threads 是传给原生的线程提示，0 表示由 Rust 侧按可用并行度决定。

注意

durationNanos 由 Java 侧两次 System.nanoTime() 测得，包含 JNI 进入与原生执行的全部时间。

注意

事件未启用时 commit 立即返回、不写任何记录；常态下每次队列执行仍有一次事件对象构造与两次时间戳读取，对象不逃逸时可被 JIT 标量替换。

参见：basalt.nativebridge.CommandQueueBuilder, basalt.torch.Graph

字段

commands()¶

本次队列包含的指令条数

注意

与一次前向中发起的算子调用数一致，不含被 Java 侧提前拒绝的调用。

threads()¶

传给原生的线程数提示

注意

取自 NativeContext.threads()，0 表示交由 Rust 侧选择。

bytes()¶

指令缓冲的字节数

注意

恒为 commands 乘 128；缓冲在多次执行间复用并按需扩容。

durationNanos()¶

本次执行的墙钟耗时，单位纳秒

注意

覆盖 JNI 调用与原生计算，不包含队列构建时间。

final class OperatorExecutionEvent extends Event¶

单次算子或估计器执行的 JFR 事件

BasaltMetrics.recordOperator 与 BasaltMetrics.time 在记录一次算子耗时后填充并 commit 该事件，带出算子名称、处理元素数、字节数与纳秒耗时；经典模型 fit 与批量推理复用同一条通道。

JFR 是 JDK 内置的事件框架，可用 jcmd <pid> JFR.start 在运行时打开采集，无需改动代码或重启进程。默认关闭时，一次记录只剩一次对象分配与一次 shouldCommit() 判断，对象不逃逸时可被 JIT 标量替换，字段写入与 commit 全部省略。

注意

事件名为 basalt.OperatorExecution，标签 Basalt Operator Execution，类别 {Basalt, Observability}。

注意

operator 为 null 或空白时整条记录提前返回：既不进统计表，也不产生事件。

注意

durationNanos 写入前被钳到至少 1 纳秒，JFR 记录里不会出现 0 耗时，统计表也用同一钳制值。

注意

elements 与 bytes 原值写入事件，而统计表用 Math.max(0, ·) 抹平负值，异常入参下两者可能不一致。

注意

是否抽样、堆栈深度与阈值都由 JFR 配置控制，事件本身不做采样。

注意

每次记录都先问 shouldCommit() 再赋值，未启用采集时不写任何字段。

参见：basalt.BasaltMetrics, basalt.torch.NativeMemoryEvent, basalt.torch.NativeQueueEvent

字段

operator()¶

算子或估计器的类别名称

注意

调用方传入的自由字符串，如 matmul、fit；null 或空白会被直接丢弃。

elements()¶

本次处理的元素或样本数

注意

原值写入事件；统计表对负值按 0 计。

bytes()¶

本次处理的字节数

注意

原值写入事件；统计表对负值按 0 计。

durationNanos()¶

本次执行的墙钟耗时，单位纳秒

注意

由调用方以 System.nanoTime() 测得，写入前钳到至少 1。

注意

覆盖 JNI 进入与原生计算，不含事件提交本身的开销。

abstract class Optimizer¶

参数列表上的优化器基类：校验、清零与全局裁剪

构造时要求参数列表非空、学习率有限且为正，列表被复制为不可变视图，之后修改原列表不影响优化器。子类只需实现 step()，在其中对每个参数调用张量级的更新原语。

zeroGrad 遍历参数清零梯度缓冲；clipGradNorm 把全部参数的梯度拼成一个向量求 L2 范数，超过阈值时按 maxNorm/(norm+1e-6) 统一缩放。

注意

learningRate 是 protected 可变字段（非 final）：构造时要求 >0，之后可由 setLearningRate 或 LRScheduler 改为含 0 的非负值。

注意

clipGradNorm 以第一个参数的上下文执行，参数跨上下文会抛 IllegalArgumentException；maxNorm 为负或非有限由原生层拒绝。

注意

step() 不清梯度，必须显式调用 zeroGrad()，否则下一批梯度会与旧梯度叠加。

注意

基类不持有计算图，也不负责清图；清图由调用方在步末完成。

参见：basalt.torch.SGD, basalt.torch.AdamW, basalt.torch.Tensor#zeroGrad()

构造器

Optimizer(List<Tensor> parameters, float learningRate)¶

校验并保存参数列表与学习率

注意

受保护构造器，仅供子类调用；不校验参数是否 requiresGrad。

参数

• parameters (List<Tensor>) – 待优化的张量列表，必须非空且不含 null 元素

• learningRate (float) – 步长，必须有限且大于 0

抛出

• IllegalArgumentException – 学习率非正或非有限

• BasaltRuntimeException – parameters 为 null

方法

parameters() → List<Tensor>¶

返回不可变的参数列表

返回

构造时保存的只读列表

learningRate() → float¶

返回当前学习率

注意

与构造入参可能不同：StepLR 等调度器每步都会回写。

返回

learningRate 字段的实时值

setLearningRate(float learningRate)¶

改写学习率，允许设为 0

直接覆盖 learningRate 字段，不重建优化器，也不触碰梯度或优化器状态。

注意

构造器要求严格大于 0，本方法放宽到 0，因此可把参数冻住。

注意

字段为普通 float，与 step() 并发调用时没有同步保证。

参数

• learningRate (float) – 新学习率，必须有限且非负；0 合法，表示冻结更新

抛出

• IllegalArgumentException – 为负或非有限（NaN、±Inf）

zeroGrad()¶

清零全部参数的梯度

注意

只清梯度缓冲，不改参数数据，也不清计算图。

step()abstract¶

按具体算法更新一次参数

注意

抽象方法；实现应只读取当前梯度并原地更新参数数据。

stepAccumulated(int accumulationSteps) → boolean¶

按累计调用次数决定是否执行更新

内部微批计数器加一；达到 accumulationSteps 时依次调用 step()、zeroGrad() 并把计数器归零。

注意

计数器是实例级 int 字段，所有参数共享，且不做同步，多线程同时调用会丢计数。

注意

更新分支顺带清梯度，外层不要再调 zeroGrad()。

注意

accumulationSteps 可逐次改变，比较对象始终是当次传入值。

参数

• accumulationSteps (int) – 每多少次调用做一次更新，必须为正；每次调用都与当前计数器比较

返回

本次触发更新返回 true，否则返回 false

抛出

• IllegalArgumentException – accumulationSteps ≤ 0

clipGradNorm(float maxNorm) → double¶

按全局 L2 范数裁剪梯度

先对所有参数的梯度求平方和再开方得到全局范数；若范数超过 maxNorm 且大于 0，则按 maxNorm/(norm+1e-6) 统一缩放每个梯度。

注意

maxNorm 非法时由原生层抛 IllegalStateException。

注意

必须在 step() 之前调用，否则缩放的是已经消费过的梯度。

参数

• maxNorm (float) – 允许的最大全局范数，必须有限且非负

返回

裁剪前的全局范数，便于监控梯度爆炸

抛出

• IllegalArgumentException – 参数张量属于不同上下文

clipGradValue(float clipValue) → double¶

以给定上界做全局范数裁剪

校验 clipValue 为正的有限值后直接委托 clipGradNorm(clipValue)：按全部参数梯度的全局 L2 范数与它比较，超限时统一缩放。

注意

方法名暗示逐元素截断，当前实现是全局范数裁剪，不会把单个梯度夹到 [-clipValue, clipValue]。

注意

与 clipGradNorm 一样须在 step() 之前调用，否则缩放的是已消费的梯度。

注意

参数跨上下文时由 clipGradNorm 抛 IllegalArgumentException。

参数

• clipValue (float) – 裁剪上界，必须有限且大于 0

返回

裁剪前的全局 L2 范数

抛出

• IllegalArgumentException – clipValue ≤ 0 或非有限

字段

parameters()¶

受优化的参数列表，构造时复制为不可变视图

注意

顺序即 step 与 zeroGrad 的遍历顺序；子类可直接读取。

learningRate()¶

全局学习率

注意

protected 可变字段且非 final；setLearningRate 与 LRScheduler 都会直接改写它。

final class SGD extends Optimizer¶

无动量、解耦权重衰减的随机梯度下降

每步执行 p -= lr * (g + wd * p)。没有动量、Nesterov 与内建裁剪；需要梯度裁剪时先调用 Optimizer.clipGradNorm。

更新在原生侧原地完成，不分配新张量、不写计算图，因此 step 本身不会增加图节点；参数数据被直接改写，已有张量引用无需替换。

注意

单参构造器把 weightDecay 设为 0。

注意

构造器不校验 weightDecay 取值，负值等价于鼓励参数增长。

注意

梯度为零时参数仍会因权重衰减项而收缩。

参见：basalt.torch.Optimizer, basalt.torch.AdamW, basalt.torch.Tensor#sgd(float, float)

构造器

SGD(List<Tensor> parameters, float learningRate)¶

构造无权重衰减的 SGD

等价于 weightDecay=0 的完整构造器。

参数

• parameters (List<Tensor>) – 待优化的张量列表，必须非空

• learningRate (float) – 步长，必须有限且大于 0

抛出

• IllegalArgumentException – 参数列表为空，或学习率非正、非有限

SGD(List<Tensor> parameters, float learningRate, float weightDecay)¶

构造带解耦权重衰减的 SGD

参数

• parameters (List<Tensor>) – 待优化的张量列表，必须非空

• learningRate (float) – 步长，必须有限且大于 0

• weightDecay (float) – 权重衰减系数，按 p -= lr*(g + wd*p) 生效

抛出

• IllegalArgumentException – 参数列表为空，或学习率非正、非有限

方法

step()¶

对全部参数执行一步 SGD 更新

注意

按参数列表顺序逐个原地更新；不清梯度、不推进步数计数。

final class SelfAttention extends Module<BSF, BSF>¶

单头自注意力，输入 [batch, time, channels]

构造四个 dimension 到 dimension 的 Linear（Q、K、V 与输出投影，均带偏置），缩放系数是值为 1/sqrt(dimension) 的标量张量；不拆分注意力头。

forward 计算 scores = Q · Kᵀ · scale，形状 [B, T, T]；causal 为真时把 column > row 的位置加上 -1e9，再沿最后一维 softmax，与 V 相乘后经输出投影返回 [B, T, C]。

注意

输入必须是三维，否则抛 IllegalArgumentException；dimension 不要求能被任何数整除。

注意

因果掩码每次前向按 time×time 重新构造；模块内没有 dropout，因此没有 train()/eval() 重写。

注意

namedParameters 的名称前缀为 query.、key.、value.、output.。

参见：basalt.torch.MultiheadAttention, basalt.torch.TransformerBlock

构造器

SelfAttention(NativeContext context, int dimension, boolean causal, long seed)¶

构造单头自注意力

注意

参数总量为 4·(dimension² + dimension)，无位置编码。

参数

• context (NativeContext) – 持有原生上下文，四个投影与缩放张量都在其中

• dimension (int) – 模型维度，也是输入最后一维与各投影的尺寸

• causal (boolean) – 是否屏蔽未来位置，为 true 时位置 t 只能看到不晚于 t 的位置

• seed (long) – 种子基数，四个投影分别偏移 101、211、307、401

方法

forward(Tensor input) → Tensor¶

计算单头自注意力输出

注意

通道数与 dimension 不一致时，在 Q 投影的 Linear 内因最后一维不匹配而抛 IllegalArgumentException。

参数

• input (Tensor) – 形状 [batch, time, channels]，channels 需等于 dimension

返回

[batch, time, channels] 的新张量

抛出

• IllegalArgumentException – 输入不是三维张量

parameters() → List<Tensor>¶

返回四个投影的参数

返回

query、key、value、output 的权重与偏置依次拼接

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回带模块前缀的命名参数

返回

名称形如 query.weight、key.bias、output.weight 等

close()¶

释放子模块与缩放张量

注意

四个投影与缩放张量都由本模块持有并一并释放。

final class Sequential<I extends Layout, O extends Layout> extends Module<I, O>¶

类型态模块流水线：add 返回新实例并推进布局类型

start() 产生空的 Sequential<L, L>，L 由目标类型或上下文推断；add 要求实参是 Module<O, N>，返回 Sequential<I, N>，于是链条上每一步都把上一层输出布局钉成下一层输入布局。

因此 Conv2d（BCHW→BCHW）之后直接接 Linear（BC→BC）无法通过编译，必须插入 Flatten（BCHW→BC）把布局改写成 BC 才能继续。

实例不可变：add 复制模块列表并返回新对象，原对象保持为空，只有最终对象代表完整流水线。forward 按插入顺序依次调用，parameters 与 namedParameters 按同样顺序拼接。

示例

import basalt.Basalt;
import basalt.core.BC;
import basalt.nn.Activation;
import basalt.torch.AdamW;
import basalt.torch.Loss;
import basalt.torch.Optimizer;
import basalt.torch.Sequential;
import basalt.torch.Tensor;

try (Basalt jsl = Basalt.create()) {
    // start() 产生空流水线，输入与输出都钉在 BC
    Sequential<BC, BC> model = Sequential.<BC>start()
            .add(basalt.linear(8, 16))
            .add(basalt.activation(Activation.RELU))
            .add(basalt.linear(16, 4));

    Tensor features = basalt.randn(32, 8);
    Tensor targets = basalt.zeros(32, 4);
    Loss criterion = basalt.mse();
    Optimizer optimizer = new AdamW(model.parameters(), 0.01f);

    for (int step = 0; step < 5; step++) {
        Tensor loss = criterion.loss(model.forward(features), targets);
        loss.backward();       // 1. 反向：梯度累积到参数
        optimizer.step();      // 2. 更新：消费当前梯度
        optimizer.zeroGrad();  // 3. 清零：为下一批准备
        basalt.clearGraph();      // 4. 清图：丢弃本步节点与中间张量
    }
}

注意

空流水线的 forward 直接返回输入引用，不做复制。

注意

namedParameters 的键带 layer{i}. 前缀并与子模块自身的键拼接，可直接用于 stateDict。

注意

train()/eval() 递归传播到子模块，但 isTraining() 只反映 Sequential 自身的标志。

注意

close() 递归关闭子模块；同一模块实例加入多条流水线时，关闭其中一条会释放共享参数。

注意

depth() 返回直接子模块个数，不是嵌套深度。

参见：basalt.torch.Module, basalt.torch.Linear, basalt.torch.Conv2d, basalt.torch.Flatten

方法

<L extends Layout> start() → Sequential<L, L>static¶

创建空流水线，输入输出布局相同

静态工厂方法，返回不含任何模块的 Sequential<L, L>；L 由赋值目标或后续 add 的实参推断。

注意

返回对象不可变，添加模块必须接收 add 的返回值。

返回

空的 Sequential，输入与输出布局均为 L

<N extends Layout> add(Module<O, N> module) → Sequential<I, N>¶

追加模块并返回新的流水线

实参类型必须是 Module<O, N>，即其输入布局等于当前流水线的输出布局；返回类型 Sequential<I, N> 把输出布局推进为 N。

注意

不修改 this，也不校验运行期形状；布局约束完全在编译期生效。

注意

同一模块实例可被追加到多条流水线，参数随之共享。

参数

• module (Module<O, N>) – 待追加的模块，输入布局必须与当前输出布局 O 一致

返回

包含新模块的新 Sequential 实例，原实例不变

forward(Tensor input) → Tensor¶

按顺序执行全部子模块

注意

逐个调用子模块的 forward 而非 apply，因此不重复做空值检查。

参数

• input (Tensor) – 流水线输入张量

返回

最后一个子模块的输出；空流水线时返回 input 本身

parameters() → List<Tensor>¶

拼接全部子模块的参数

注意

每次调用都重新构建列表，热路径上不宜反复调用。

返回

按模块插入顺序、模块内顺序展开的新列表

namedParameters() → List<Map.Entry<String, Tensor>>¶

拼接带层前缀的参数名

注意

前缀用点号连接，与子模块自身的命名约定拼接。

返回

键形如 layer0.weight、layer2.bias 的条目列表

train()¶

切换到训练态并递归传播

注意

先设置自身标志，再对每个子模块调用 train()。

eval()¶

切换到推理态并递归传播

注意

先设置自身标志，再对每个子模块调用 eval()；之后 Dropout 直通、BatchNorm2d 使用运行统计。

close()¶

递归关闭全部子模块

注意

不关闭外部传入的张量；共享子模块被多条流水线持有时会被关闭多次，句柄关闭幂等。

depth() → int¶

返回流水线中的模块个数

注意

嵌套的 Sequential 只计为一层。

返回

直接子模块数量

modules() → List<Module<?, ?>>¶

返回子模块的只读列表

注意

与 depth() 用的是同一列表，两者长度恒等。

注意

列表不可修改，追加模块只能通过 add 取得新实例。

返回

构造时 List.copyOf 保存的不可变列表，顺序即 forward 顺序

final class SmoothL1Loss implements Loss¶

带 beta 转折点的 Smooth L1（Huber）损失

前向对每个元素按 d = prediction − target 与 beta 的关系分段：|d| < beta 取 d²/(2β)，否则取 |d| − β/2；最后对所有元素取算术平均，返回形状 [1] 的标量张量。

实现全由公开算子组合：差值乘上逐元素构造的系数与偏移张量，再与 1/n 的全 1 向量做矩阵乘求均值，因此反向梯度与解析式一致（内侧 d/β，外侧 sign(d)/n）。

beta ≤ 0 时直接退化为 l1Loss，即平均绝对误差。

注意

默认构造器取 beta = 1.0，与 torch.nn.SmoothL1Loss 的默认分段一致。

注意

构造器只拒绝负 beta，不校验有限性；NaN 会一路传播到输出。

注意

prediction 与 target 必须逐维相等，否则抛 IllegalArgumentException。

注意

均值除以元素总数而非 batch 大小；分段掩码在 Java 侧逐元素生成，每次前向多一次 data() 全量拷贝。

参见：basalt.torch.Tensor#smoothL1Loss(Tensor, float), basalt.torch.L1Loss, basalt.torch.MSELoss

构造器

SmoothL1Loss()¶

以 beta=1 构造

等价于 SmoothL1Loss(1.0f)。

SmoothL1Loss(float beta)¶

指定转折点构造损失

注意

beta 越大，平方段越宽、对离群点越敏感。

参数

• beta (float) – 分段转折点，必须非负；0 使损失退化为 L1，非有限值不被拒绝

抛出

• IllegalArgumentException – beta 为负（消息为 beta must be non-negative）

方法

beta() → float¶

返回转折点

返回

构造时保存的 beta 值

loss(Tensor prediction, Tensor target) → Tensor¶

计算逐元素平均的 Smooth L1

注意

形状校验在 beta 分支之前完成，beta ≤ 0 时转调 l1Loss 也不会重查形状。

注意

不修改输入，也不引入可训练参数。

参数

• prediction (Tensor) – 预测张量，形状须与 target 完全相同

• target (Tensor) – 目标张量，与 prediction 逐元素作差

返回

形状 [1] 的标量损失张量

抛出

• IllegalArgumentException – 两侧形状不完全相同（消息为 SmoothL1Loss expects equal shapes）

final class StepLR implements LRScheduler¶

每隔 stepSize 步按 gamma 衰减学习率的调度器

构造时把优化器当前学习率定格为基准 baseLr，并保存 stepSize 与 gamma；每次 step() 先把内部计数器加一，再用整数除法 decays = currentStep / stepSize 求衰减次数，写回 baseLr · gamma^decays。

于是从 0 开始的前 stepSize−1 次调用学习率保持 baseLr，第 stepSize 次调用恰好衰减一次，之后每再走 stepSize 步再乘一次 gamma；学习率按 baseLr·gamma^k 阶梯下降，不会逐次累乘产生漂移。

调度器只改写优化器的学习率字段，不触碰梯度、动量与计算图，其节奏必须与优化器 step() 的调用对齐。

注意

stepSize ≤ 0 抛 IllegalArgumentException，gamma ≤ 0 或非有限同样抛 IllegalArgumentException，optimizer 为 null 抛 NullPointerException。

注意

gamma 未限制小于 1，大于 1 时学习率阶梯上升。

注意

baseLr 构造时定格：之后手动 setLearningRate 或换用别的调度器都不改变已算好的基准。

注意

计数器是普通 int 且不参与同步，多线程共用同一调度器会重复计数。

注意

与 torch.optim.lr_scheduler.StepLR 对齐：第 e 步的学习率为 baseLr · gamma^(e // stepSize)。

参见：basalt.torch.LRScheduler, basalt.torch.Optimizer#setLearningRate(float), basalt.torch.CosineAnnealingLR

构造器

StepLR(Optimizer optimizer, int stepSize, float gamma)¶

绑定优化器并记录基准学习率

注意

构造不推进计数器，优化器学习率保持原值直到首次 step()。

参数

• optimizer (Optimizer) – 被调度的优化器，非 null，其当前学习率作为 baseLr

• stepSize (int) – 衰减间隔，必须为正；单位由调用方决定（步或 epoch）

• gamma (float) – 每次衰减的乘数，必须有限且大于 0

抛出

• NullPointerException – optimizer 为 null

• IllegalArgumentException – stepSize ≤ 0，或 gamma ≤ 0、非有限

方法

step() → float¶

推进一步并写回新的学习率

计数器自增后按 currentStep / stepSize 求已发生的衰减次数，把 baseLr · gamma^decays 写回优化器。

注意

指数以 double 调 Math.pow 计算后强转为 float。

注意

与 epoch 无绑定，调用次数即步数，语义由调用方决定。

注意

即使数值未变也会写回一次。

返回

写回后的学习率

currentLearningRate() → float¶

返回优化器当前的学习率

注意

直接读字段而不重算公式，因此能反映外部对 setLearningRate 的改动。

返回

optimizer.learningRate() 的实时值

currentStep() → int¶

返回已推进的步数

注意

它不等于衰减次数；衰减次数需用 currentStep / stepSize 求得。

返回

内部计数器，初始为 0，每次 step() 加一

final class Tensor implements AutoCloseable¶

Rust 内核支撑的 eager 张量：即时求值、反向求导与参数句柄

每次算子调用都立即在原生引擎执行并返回新张量，同时在计算图中登记一条节点，因此前向、损失与反向可以按语句顺序逐步执行，无需先构图再运行。用 parameter 创建的张量在 backward() 后累积梯度。

形状在运行期动态决定，与 torch.Tensor 一致；编译期的布局安全由上层的 Module/Sequential 泛型提供。Java 对象只保存不透明句柄与形状副本，数据始终留在原生连续缓冲中。

示例

import basalt.Basalt;
import basalt.torch.Tensor;

try (Basalt jsl = Basalt.create()) {
    Tensor features = basalt.randn(4, 3);          // 形状 [4, 3]，不求导
    Tensor weight = basalt.parameter(3, 1);        // 可训练参数，零初始化
    Tensor loss = features.matMul(weight)       // [4, 3] x [3, 1] -> [4, 1]
            .mse(basalt.zeros(4, 1));              // 标量损失，形状 [1]

    loss.backward();                            // 梯度累加到 weight
    float[] gradient = weight.grad();            // 复制出梯度
    weight.sgd(0.1f, 0.0f);                     // weight -= lr * grad
    weight.zeroGrad();                          // 清梯度
    basalt.clearGraph();                           // 清图，避免节点累积
}

注意

算子只登记到所属上下文的当前计算图：同一上下文内连续训练需在每步末清图，否则节点与中间张量持续累积。

注意

关闭任何被计算图引用的张量会清空整张图，因此中间张量不要在 backward() 之前 close()。

注意

backward() 既不自动清图也不清零梯度；梯度按累加写入，重复调用会叠加。

注意

算子输出张量的 requiresGrad 为 false，但仍作为图节点参与反向，梯度只累加到 requiresGrad 为 true 的张量上。

注意

data()/grad() 每次都把原生缓冲复制成新 float[]，热路径上频繁调用会产生额外拷贝与 GC 压力。

注意

张量绑定创建它的 NativeContext，跨上下文混用会在队列执行时被拒绝；单个张量的元素总数上限为 Integer.MAX_VALUE。

参见：basalt.torch.Module, basalt.torch.Graph, basalt.torch.Optimizer, basalt.nativebridge.NativeContext

方法

zeros(NativeContext context, int... shape) → Tensorstatic¶

创建零初始化、不参与求导的张量

注意

也用于承接算子输出，因此中间结果的 requiresGrad 为 false。

参数

• context (NativeContext) – 张量所属的原生上下文，之后所有算子必须使用同一上下文

• shape (int...) – 各维长度，必须为正；元素总数受 Java float[] 长度上限约束

返回

形状为 shape、requiresGrad 为 false 的新张量

抛出

• IllegalArgumentException – shape 为空或含非正维度

• ArithmeticException – 元素总数溢出 int

of(NativeContext context, float[] data, int... shape) → Tensorstatic¶

由展平数据与形状创建张量，默认不求导

注意

原生侧持有副本，之后修改入参数组不影响张量。

参数

• context (NativeContext) – 原生上下文

• data (float[]) – 按行优先展平的连续数据，长度必须等于 shape 的元素总数

• shape (int...) – 目标形状

返回

数据已复制进原生缓冲的新张量

抛出

• IllegalArgumentException – data 长度与 shape 元素数不一致

of(NativeContext context, float[] data, boolean requiresGrad, int... shape) → Tensorstatic¶

由数据创建张量并显式指定是否需要梯度

注意

给外部数据赋非零初值时使用此重载，避免再调用 setData。

参数

• context (NativeContext) – 原生上下文

• data (float[]) – 按行优先展平的数据，长度必须等于 shape 元素数

• requiresGrad (boolean) – true 时该张量在反向中累积梯度

• shape (int...) – 目标形状

返回

新张量

抛出

• IllegalArgumentException – data 长度与 shape 元素数不一致

parameter(NativeContext context, int... shape) → Tensorstatic¶

创建零初始化的可训练参数

注意

初始化恒为零，需要非零初值时应调用 randomize 或改用带 requiresGrad 的 of。

参数

• context (NativeContext) – 原生上下文

• shape (int...) – 参数形状

返回

requiresGrad 为 true 的零张量，可直接交给 Optimizer

random(NativeContext context, int[] shape, long seed, float scale, boolean requiresGrad) → Tensorstatic¶

以给定种子生成均匀分布张量

注意

分布是均匀分布而非正态分布；同 seed、同形状、同 scale 的结果完全可复现。

注意

scale 非法时由原生层抛 IllegalStateException。

参数

• context (NativeContext) – 原生上下文

• shape (int[]) – 形状数组，注意此处是 int[] 而非可变参数

• seed (long) – xorshift 状态种子；0 会被替换为 1，因此 0 与 1 产生同一序列

• scale (float) – 均匀分布半径，取值落在 [-scale, scale]，必须有限且非负

• requiresGrad (boolean) – true 时作为可训练参数参与反向

返回

已填充随机值的新张量

scalar(NativeContext context, float value) → Tensorstatic¶

创建形状 [1]、含单个元素的张量

注意

形状是 [1] 而不是零维，因此 rank() 为 1、size() 为 1。

参数

• context (NativeContext) – 原生上下文

• value (float) – 唯一元素的值

返回

形状为 [1] 的新张量

context() → NativeContext¶

返回张量所属的原生上下文

返回

创建该张量的 NativeContext

handle() → NativeTensorHandle¶

返回底层原生张量句柄

注意

绕过句柄直接调用会跳过 Java 侧的形状校验；张量 close() 后句柄不可再用。

返回

不透明的 NativeTensorHandle

shape() → int[]¶

返回形状副本

返回

新数组，修改它不影响张量

rank() → int¶

返回秩，即维度个数

返回

shape() 的长度

size() → int¶

返回元素总数

返回

各维长度的乘积

requiresGrad() → boolean¶

返回该张量是否需要梯度

注意

标志在构造时固定，没有原地切换方法。

返回

创建时确定的标志；只有 true 的张量会接收反向累积的梯度

data() → float[]¶

把张量数据复制为新的 float 数组

注意

每次调用都跨 JNI 复制整个缓冲；返回值可自由修改。

返回

长度等于 size() 的新数组，按行优先排列

grad() → float[]¶

把累积梯度复制为新的 float 数组

注意

梯度按累加写入，直到 zeroGrad() 或优化器清零。

返回

长度等于 size() 的新数组；从未参与反向时为全零

setData(float[] values)¶

用给定数组原地覆盖张量数据

注意

直接写原生缓冲，不产生计算图节点；常用于加载权重。

参数

• values (float[]) – 长度必须等于 size()，按行优先解释

抛出

• IllegalArgumentException – values 长度与张量元素数不一致

• BasaltRuntimeException – values 为 null

fill(float value)¶

用常数原地填充张量

注意

非有限值由原生层拒绝，抛 IllegalStateException。

参数

• value (float) – 填充值，必须有限

randomize(long seed, float scale)¶

原地重新随机化张量数据

注意

与 Tensor.random 使用同一生成器与同一区间 [-scale, scale]，常用来给 parameter 赋初值。

参数

• seed (long) – xorshift 种子；0 会被替换为 1

• scale (float) – 均匀分布半径，必须有限且非负

zeroGrad()¶

把梯度缓冲清零

注意

只清梯度，不改数据，也不清计算图。

backward()¶

从当前张量反向传播并累积梯度

以全 1 向量作为初始余切，沿计算图逆拓扑序传播，把梯度累加到所有 requiresGrad 为 true 的祖先张量上。

注意

对标量损失等价于 d loss/d loss = 1；对非标量根则是全 1 加权的向量-雅可比积。

注意

根张量必须是当前计算图的产物，否则原生层抛 IllegalStateException；清图或关闭中间张量后再调用会失败。

注意

不清图、不清零梯度，训练循环需自行安排清图与清零。

抛出

• IllegalArgumentException – 本张量属于另一个上下文

sgd(float learningRate, float weightDecay)¶

对当前张量原地执行一步 SGD 更新

注意

不做参数校验，也不清零梯度；通常经 Optimizer.step() 间接调用。

参数

• learningRate (float) – 步长，正值表示沿负梯度方向

• weightDecay (float) – 解耦权重衰减系数，按 p -= lr*(g + wd*p) 生效

adam(float learningRate, float beta1, float beta2, float epsilon, float weightDecay, long step)¶

对当前张量原地执行一步 AdamW 更新

注意

一阶、二阶矩缓冲按张量地址存放在上下文中，首次调用时惰性分配；step 必须单调递增，否则偏差校正失真。

参数

• learningRate (float) – 步长

• beta1 (float) – 一阶矩的指数衰减率

• beta2 (float) – 二阶矩的指数衰减率

• epsilon (float) – 分母稳定项

• weightDecay (float) – 解耦权重衰减系数

• step (long) – 从 1 开始的全局步数，用于偏差校正；为 0 时直接返回、不更新

norm() → double¶

返回数据向量的 L2 范数

注意

与形状无关，可用于监控参数规模。

返回

全部元素平方和的平方根，以 double 返回

gradNorm() → double¶

返回梯度向量的 L2 范数

注意

这是单张量范数；跨参数的全局裁剪请用 Optimizer.clipGradNorm。

返回

梯度平方和的平方根；梯度全零时返回 0

isFinite() → boolean¶

判断数据是否全部为有限值

注意

只检查数据，不检查梯度；训练发散时常配合 gradNorm 一起观察。

返回

所有元素既非 NaN 也非 ±Inf 时为 true

add(Tensor other) → Tensor¶

逐元素加法，支持广播

注意

Java 侧先算出输出形状，原生层再复核；结果 requiresGrad 为 false，但作为图节点参与反向。

参数

• other (Tensor) – 右操作数，形状需能与本张量右对齐广播：逐维相同、或其中之一为 1、或该维缺失

返回

形状为广播结果的新张量

抛出

• IllegalArgumentException – 两个形状无法广播

sub(Tensor other) → Tensor¶

逐元素减法，支持广播

参数

• other (Tensor) – 减数，广播规则与 add 相同

返回

形状为广播结果的新张量

抛出

• IllegalArgumentException – 两个形状无法广播

mul(Tensor other) → Tensor¶

逐元素乘法，支持广播

注意

反向时两侧梯度互为对方的输入值，广播维的梯度会沿该维求和。

参数

• other (Tensor) – 右操作数，广播规则与 add 相同

返回

形状为广播结果的新张量

抛出

• IllegalArgumentException – 两个形状无法广播

matMul(Tensor other) → Tensor¶

二维矩阵乘法 [M,K]×[K,N]

注意

只接受秩 2，不隐式升维或降维；批量情形用 batchedMatMul。

参数

• other (Tensor) – 右操作数，必须秩 2，且第一维等于本张量的第二维

返回

形状 [M,N] 的新张量

抛出

• IllegalArgumentException – 任一侧秩不为 2，或 left[1] != right[0]

batchedMatMul(Tensor other) → Tensor¶

批量矩阵乘法 [B,M,K]×[B,K,N]

注意

batch 维不做广播，必须逐元素相同。

参数

• other (Tensor) – 秩 3 的右操作数，B 与 K 必须与本张量一致

返回

形状 [B,M,N] 的新张量

抛出

• IllegalArgumentException – 秩不为 3、batch 不一致，或 left[2] != right[1]

relu() → Tensor¶

逐元素 ReLU，形状不变

注意

0 处的导数取 0，因此恰好为 0 的输入不会获得梯度。

返回

同形状新张量，负值置零

sigmoid() → Tensor¶

逐元素 sigmoid，形状不变

注意

实现按符号分支计算指数，避免大负数输入时溢出。

返回

同形状新张量，取值落在 (0,1)

tanh() → Tensor¶

逐元素双曲正切，形状不变

返回

同形状新张量，取值落在 (-1,1)

gelu() → Tensor¶

逐元素 GELU，形状不变

注意

使用含 0.044715 三次项的 tanh 近似，与 erf 精确版本存在微小差异。

返回

同形状新张量

silu() → Tensor¶

逐元素 SiLU，形状不变

注意

定义为 x * sigmoid(x)，也称 swish。

返回

同形状新张量

softmax(int dimension) → Tensor¶

沿最后一维做数值稳定的 softmax

注意

任意秩都可使用，前导维被视作彼此独立的行。

注意

内部先减行内最大值再取指数，避免上溢；分母非有限时原生层报错。

参数

• dimension (int) – 维索引；负数先加 rank 归一化；归一化后必须等于 rank-1

返回

同形状新张量，最后一维每行之和为 1

抛出

• IllegalArgumentException – 归一化后仍不是最后一维（消息为 softmax currently supports the last dimension）

crossEntropy(Tensor targets) → Tensor¶

softmax 与负对数似然的融合损失

注意

目标越界或非有限由原生层抛 IllegalStateException；targets 的秩不限，只校验元素个数。

注意

前向用 log-sum-exp 稳定计算，反向一次得到 (softmax - onehot)/N。

参数

• targets (Tensor) – 元素个数必须等于 N 的目标张量，元素为 [0,C) 的类索引，以 float 存储

返回

形状 [1] 的标量张量，按 N 取平均

抛出

• IllegalArgumentException – logits 秩不为 2，或 targets.size() != N

mse(Tensor target) → Tensor¶

逐元素平方误差的平均值

注意

除以元素总数而非 batch 大小；反向为 2*(p-t)/n，对两侧同时累积。

参数

• target (Tensor) – 目标张量，形状必须与本张量完全相同

返回

形状 [1] 的标量张量

抛出

• IllegalArgumentException – 两个形状不完全相同（消息为 MSE expects equal shapes）

layerNorm(Tensor gamma, Tensor beta, float epsilon) → Tensor¶

沿最后一维做层归一化

注意

方差按该维元素数取有偏估计；前导维各自独立归一化。

注意

Java 侧只校验 gamma 与 beta 的元素数，原生层还要求两者形状恰为 [dim]，不符时抛 IllegalStateException。

参数

• gamma (Tensor) – 缩放参数，元素数必须等于最后一维长度

• beta (Tensor) – 平移参数，元素数必须等于最后一维长度

• epsilon (float) – 加在方差上的稳定项，避免除零

返回

同形状新张量

抛出

• IllegalArgumentException – gamma 或 beta 的元素数不等于最后一维长度

batchNorm2d(Tensor gamma, Tensor beta, Tensor runningMean, Tensor runningVar, float epsilon, float momentum, boolean training) → Tensor¶

NCHW 输入的批归一化

注意

归一化使用有偏方差，更新 runningVar 时改用无偏估计（spatial>1 时乘 n/(n-1)）。

注意

runningMean/runningVar 必须由调用方持有并跨步复用，否则推理阶段每次都从初值开始。

注意

原生层要求 gamma、beta、runningMean、runningVar 的形状恰为 [C]，不符时抛 IllegalStateException。

参数

• gamma (Tensor) – 逐通道缩放，元素数必须等于通道数 C

• beta (Tensor) – 逐通道平移，元素数必须等于 C

• runningMean (Tensor) – 长度 C 的运行均值，训练态被原地更新

• runningVar (Tensor) – 长度 C 的运行方差，训练态被原地更新

• epsilon (float) – 加在方差上的稳定项

• momentum (float) – 运行统计的滑动系数，new = (1-m)*old + m*batch

• training (boolean) – true 用批统计并更新运行统计，false 用运行统计做推理

返回

同形状新张量

抛出

• IllegalArgumentException – 输入秩不为 4（消息为 batchNorm2d expects NCHW input）

conv2d(Tensor weight, int kernelHeight, int kernelWidth, int strideHeight, int strideWidth, int paddingHeight, int paddingWidth) → Tensor¶

NCHW 输入与 OIHW 权重的二维卷积

注意

Hout = (H + 2*ph - kh)/sh + 1，整除截断，不做向上取整。

注意

stride 为 0 时 Java 侧整数除法先抛 ArithmeticException。

注意

kernelHeight/kernelWidth 只用于推算输出尺寸，原生层以 weight 的后两维为准；两者不一致会因输出形状不匹配被拒绝。

注意

实现为 im2col 加 GEMM，不包含偏置，需要时对输出调用 add。

参数

• weight (Tensor) – 形状 [O, I, kh, kw] 的卷积核，I 必须等于输入通道数

• kernelHeight (int) – 卷积核高，必须与 weight 的第三维一致

• kernelWidth (int) – 卷积核宽，必须与 weight 的第四维一致

• strideHeight (int) – 纵向步长，必须 ≥1

• strideWidth (int) – 横向步长，必须 ≥1

• paddingHeight (int) – 纵向单侧零填充量

• paddingWidth (int) – 横向单侧零填充量

返回

形状 [N, O, Hout, Wout] 的新张量

抛出

• IllegalArgumentException – 按公式算出的输出高或宽 ≤0（消息为 conv2d output would be empty）

conv2d(Tensor weight, int kernelHeight, int kernelWidth, int strideHeight, int strideWidth, int paddingHeight, int paddingWidth, int dilationHeight, int dilationWidth, int groups) → Tensor¶

带膨胀与分组的二维卷积

与七参重载共用输出尺寸推算，但有效核尺寸改为 dilation·(k−1)+1：Hout = (H + 2·ph − effKh)/sh + 1，整除截断。groups 把输入输出通道切成互不相通的分组，每组用 weight 的对应通道段独立卷积。

注意

七参重载等价于 dilation=1、groups=1 的本方法。

注意

输出尺寸向下取整，尾部不足一个窗口的部分被丢弃。

注意

只有 groups=1 且 dilation 均为 1 时才走 im2col 快速路径，其余情形走原生分组分支。

注意

weight 跨上下文时在队列执行阶段被拒绝。

参数

• weight (Tensor) – 形状 [O, I/groups, kh, kw] 的卷积核，秩必须为 4

• kernelHeight (int) – 卷积核高，参与有效核尺寸推算

• kernelWidth (int) – 卷积核宽，参与有效核尺寸推算

• strideHeight (int) – 纵向步长，须 ≥1，为 0 时整除先抛异常

• strideWidth (int) – 横向步长，须 ≥1

• paddingHeight (int) – 纵向单侧零填充量

• paddingWidth (int) – 横向单侧零填充量

• dilationHeight (int) – 纵向膨胀系数，按 dilation·(kh−1)+1 放大采样间距；Java 侧不校验取值

• dilationWidth (int) – 横向膨胀系数，同上

• groups (int) – 分组数，须为正，且输入通道与输出通道都能被它整除

返回

形状 [N, O, Hout, Wout] 的新张量，不含偏置

抛出

• IllegalArgumentException – groups ≤ 0、输入或权重秩不为 4、通道不满足整除关系，或输出高宽 ≤0

• ArithmeticException – stride 为 0

groupNorm(Tensor gamma, Tensor beta, int groups, float epsilon) → Tensor¶

按通道分组做归一化

把 C 个通道分成 groups 组，每组将 channels_per_group × spatial 个元素拉平后求均值与有偏方差，以 sqrt(var + epsilon) 为分母归一化，再逐通道乘 gamma、加 beta。统计在每个样本的每个分组内独立计算，不维护运行统计。

注意

输入按 [N, C, ...] 解释，C 之后的维度全部并入 spatial；spatial 为 0 时由原生层拒绝。

注意

方差为有偏估计，除以 group_elems 而非 group_elems−1；均值与方差在原生侧以 f64 累加。

注意

groups 为负时整除判断可能通过，最终由原生层形状校验拒绝。

注意

没有 running mean/var，训练与推理行为完全一致。

参数

• gamma (Tensor) – 逐通道缩放，元素数必须等于通道数 C

• beta (Tensor) – 逐通道平移，元素数必须等于 C

• groups (int) – 分组数，须使 C 能被整除；为 0 时取模先抛异常

• epsilon (float) – 加在方差上的稳定项，按 var + epsilon 开方

返回

与输入同形状的新张量

抛出

• IllegalArgumentException – 输入秩小于 2、C 不能被 groups 整除，或 gamma/beta 元素数不等于 C

• ArithmeticException – groups 为 0

leakyRelu(float negativeSlope) → Tensor¶

带负斜率的 ReLU

用公开算子组合实现：正部取 relu()，负部把输入取反再 relu 后乘上 negativeSlope，两半相加，等价于 x>0 ? x : negativeSlope·x。

注意

至少产生 6 个中间张量，全部留在当前计算图内直到清图。

注意

x 恰为 0 时两条分支都取 0，梯度为 0。

注意

传入 NaN 时 0×NaN 仍为 NaN，整个输出被污染。

参数

• negativeSlope (float) – 负半支斜率，常取 0.01；不做范围校验，取 1 时退化为恒等映射

返回

同形状新张量

elu(float alpha) → Tensor¶

指数线性单元，负半支按 alpha 缩放

前向在 Java 侧读回全部元素逐点构造常量张量 mask 与 offset，再以 x·mask + offset 表达结果：x>0 时 mask=1、offset=0，输出即 x；x≤0 时输出恰为 alpha·(eˣ−1)。

注意

属主机侧混合实现：一次 data() 全量拷贝、两个 float[] 临时数组与两次 of() 上传，不经过指令队列。

注意

梯度由通用图组合得出，x<0 处为 alpha·eˣ，与解析导数一致；x=0 归入负半支，导数取 alpha。

注意

每次前向都要往返一次 JNI 与主机循环，大张量上明显慢于纯原生一元算子。

注意

alpha 为 NaN 时整个输出为 NaN。

参数

• alpha (float) – 负半支饱和系数；源码不校验取值

返回

同形状新张量

l1Loss(Tensor target) → Tensor¶

逐元素绝对误差的平均值

先校验形状完全相同，再用 relu 与取反组合出 |d|（d = 本张量 − target），最后与长度 size() 的 1/n 全 1 向量做一次矩阵乘求均值。

注意

除以元素总数而非 batch 大小；梯度为 sign(x − target)/n，对两侧同时累积。

注意

由 sub、relu、mul、add、reshape、matMul 等公开算子组合，没有专门的原生 L1 核。

注意

内部 reshape 到 [1, size()] 再乘，秩超过 4 的输入会在指令队列层面被拒绝。

参数

• target (Tensor) – 目标张量，须与本张量逐维相等

返回

形状 [1] 的标量张量，等于全部元素 |x − target| 的算术平均

抛出

• IllegalArgumentException – 两侧形状不完全相同（消息为 L1Loss expects equal shapes）

smoothL1Loss(Tensor target, float beta) → Tensor¶

逐元素分段平滑 L1 的平均值

先校验形状相等；beta ≤ 0 时转调 l1Loss。beta > 0 时对每个元素取 d = 本张量 − target：|d| < beta 记 d²/(2β)，否则记 |d| − β/2，最后对全部元素取平均。

注意

分段靠两组常量张量实现：内侧 d·(1/√(2β)) 再平方，外侧 d·sign(d) − β/2，二者按元素互补。

注意

需要一次 diff.data() 全量拷贝来生成系数；beta 未校验有限性，NaN 会让输出整体变成 NaN。

注意

beta=1 时与 torch.nn.SmoothL1Loss 的默认分段一致。

注意

梯度内侧为 d/β、外侧为 sign(d)，最后统一除以元素总数。

参数

• target (Tensor) – 目标张量，须与本张量逐维相等

• beta (float) – 分段转折点，须为正才启用平方支；≤0 时退化为 L1

返回

形状 [1] 的标量张量

抛出

• IllegalArgumentException – 两侧形状不完全相同（消息为 SmoothL1Loss expects equal shapes）

bceWithLogitsLoss(Tensor target) → Tensor¶

带 logits 的二元交叉熵，按元素取平均

前向逐元素计算 max(x,0) − x·y + log1p(exp(−|x|)) 后除以元素总数，得到形状 [1] 的标量；该写法在 x 取大正值或大负值时都不会上溢。

反向不依赖 sigmoid 的自动微分链，而是把精确梯度 (sigmoid(x) − y)/n 作为常数系数挂在一个对 logits 线性的代理节点上，因此一阶导与解析式一致。

注意

入参是 logits，内部自做 sigmoid，不要再传概率。

注意

只对齐一阶梯度：代理节点不表达二阶导，也不向 target 回传梯度。

注意

每次调用会做 data()、target.data() 与 sigmoid 结果 data() 三次全量拷贝，属主机侧实现，不走原生损失核。

参数

• target (Tensor) – 目标张量，须与本张量逐维相等；元素通常取 [0,1]，源码不做范围校验

返回

形状 [1] 的标量张量，等于全部元素损失的平均

抛出

• IllegalArgumentException – 两侧形状不完全相同（消息为 BCEWithLogitsLoss expects equal shapes）

maxPool2d(int kernelHeight, int kernelWidth, int strideHeight, int strideWidth) → Tensor¶

NCHW 输入的最大池化，无填充

注意

Hout = (H - kh)/sh + 1，尾部不足一个窗口的部分被丢弃。

注意

核大于输入时 Java 公式可能仍给出正值，最终由原生层抛 IllegalStateException。

注意

stride 为 0 时抛 ArithmeticException；梯度只回传到每个窗口最大值所在位置。

参数

• kernelHeight (int) – 窗口高，必须为正

• kernelWidth (int) – 窗口宽，必须为正

• strideHeight (int) – 纵向步长，必须 ≥1

• strideWidth (int) – 横向步长，必须 ≥1

返回

形状 [N, C, Hout, Wout] 的新张量

抛出

• IllegalArgumentException – 输出高或宽 ≤0（消息为 maxPool2d output would be empty）

dropout(float probability, boolean training, long seed) → Tensor¶

训练态按概率置零并做反向缩放

注意

采用 inverted dropout，保留元素乘 1/(1-p)，推理阶段无需再缩放。

注意

掩码由 xorshift 生成，同 seed、同形状可复现。

注意

即使 training=false 也会执行一次队列调用并登记图节点。

参数

• probability (float) – 置零概率，取值 [0,1)；越界由原生层拒绝

• training (boolean) – false 时原样输出，既不置零也不缩放

• seed (long) – 掩码种子，反向会用同一种子重建同一掩码

返回

同形状新张量

embedding(Tensor weight) → Tensor¶

按索引在权重表中查表

注意

输出比索引多一维：[B,T] 索引得到 [B,T,dim]。

注意

索引以 float 存储并被截断取整，越界或非有限由原生层抛 IllegalStateException。

注意

反向对重复索引做散射累加。

参数

• weight (Tensor) – 秩 2 的 [vocab, dim] 嵌入表，通常由 parameter 创建

返回

形状为索引形状后追加 dim 的新张量

抛出

• IllegalArgumentException – weight 秩不为 2（消息为 embedding weight must be rank 2）

permute(int... permutation) → Tensor¶

按给定顺序重排维度

注意

指令队列只接受秩 1..4，更高秩会抛 IllegalArgumentException。

注意

返回新张量并实际搬运数据，不是零拷贝视图。

参数

• permutation (int...) – 长度必须等于秩，元素为 [0,rank) 的维索引

返回

形状为 input[perm[0]], input[perm[1]]… 的新张量

抛出

• IllegalArgumentException – 存在越界索引（消息为 permutation index out of range）

reshape(int... newShape) → Tensor¶

改变形状并保持元素顺序

注意

秩限制为 1..4，由指令队列校验；数据按行优先复制，与原张量不共享存储。

参数

• newShape (int...) – 各维必须为正，元素总数必须等于 size()

返回

形状为 newShape 的新张量

抛出

• IllegalArgumentException – 元素总数不一致（消息含实际数量与目标数量）

detachCopy() → Tensor¶

复制数据并断开与计算图的联系

注意

与普通算子的区别：不写图节点、不参与反向、不共享存储。

注意

修改副本不影响原张量，反之亦然；常用于记录指标或构造目标张量。

返回

requiresGrad 为 false、数据相同的新张量

close()¶

释放原生张量句柄

注意

幂等，重复调用无副作用；上下文关闭后调用也不会报错。

注意

若该张量被当前计算图引用，会连带清空整张图。

toString() → String¶

返回形状与 requiresGrad 的简短描述

注意

不含元素值，避免打印大张量时触发整块复制。

返回

形如 Tensor[2, 3][requiresGrad=true] 的字符串

final class TransformerBlock extends Module<BSF, BSF>¶

pre-norm 结构的 Transformer 块

前向为 pre-norm 顺序：先对输入做 LayerNorm 再进多头注意力，注意力输出过 dropout 后与输入相加；所得残差再做 LayerNorm、升维 Linear、GELU、降维 Linear，过 dropout 后与残差相加。

两个 LayerNorm 的 epsilon 固定为 1e-5；FFN 隐藏维由 hiddenSize 指定，激活固定为 GELU；注意力输出与 FFN 输出各过一次 dropout，两者共用同一个 Dropout 实例。

六参构造器把 heads 固定为 1，等价于单头注意力；七参构造器可显式指定头数。

注意

dimension 必须能被 heads 整除，否则构造时由 MultiheadAttention 抛 IllegalArgumentException。

注意

共用 dropout 实例意味着两次 dropout 使用相邻但不同的种子，掩码互不相同。

注意

train()/eval() 递归切换 norm1、attention、norm2、两个 Linear 与 dropout；parameters() 不含 dropout 与激活，因为它们没有参数。

注意

namedParameters 的名称前缀为 norm1.、attention.、norm2.、ffn1.、ffn2.。

参见：basalt.torch.MultiheadAttention, basalt.torch.LayerNorm, basalt.torch.Activation

构造器

TransformerBlock(NativeContext context, int dimension, int hiddenSize, boolean causal, float dropoutProbability, long seed)¶

构造单头注意力的 Transformer 块

注意

等价于 heads=1 的七参构造器。

参数

• context (NativeContext) – 持有原生上下文，全部子模块都在其中分配

• dimension (int) – 模型维度，也是输入与输出的最后一维

• hiddenSize (int) – FFN 中间层宽度，通常取 dimension 的若干倍

• causal (boolean) – 是否在注意力中使用因果掩码

• dropoutProbability (float) – 注意力权重与两处残差前的置零概率，取值 [0, 1)

• seed (long) – 种子基数，注意力偏移 11、两个 Linear 偏移 23 与 37、dropout 偏移 41

TransformerBlock(NativeContext context, int dimension, int hiddenSize, int heads, boolean causal, float dropoutProbability, long seed)¶

构造可指定头数的 Transformer 块

注意

两个 LayerNorm 使用 1e-5 作为 epsilon，FFN 激活固定为 GELU，不可配置。

参数

• context (NativeContext) – 持有原生上下文，全部子模块都在其中分配

• dimension (int) – 模型维度，必须能被 heads 整除

• hiddenSize (int) – FFN 中间层宽度

• heads (int) – 多头注意力的头数

• causal (boolean) – 是否在注意力中使用因果掩码

• dropoutProbability (float) – 注意力权重与两处残差前的置零概率，取值 [0, 1)

• seed (long) – 种子基数，各子模块按固定偏移派生

抛出

• IllegalArgumentException – dimension 不能被 heads 整除

方法

forward(Tensor input) → Tensor¶

执行注意力残差与 FFN 残差

注意

两处残差都加在未归一化的主干上，归一化只作用于子层输入，因此主干数值随层数累积。

参数

• input (Tensor) – 形状 [batch, time, channels]，channels 等于 dimension

返回

[batch, time, channels] 的新张量

抛出

• IllegalArgumentException – 输入不是三维，或通道数与 dimension 不一致

parameters() → List<Tensor>¶

返回全部子模块的可训练参数

返回

依次为 norm1、attention、norm2、ffn1、ffn2 的参数

namedParameters() → List<Map.Entry<String, Tensor>>¶

返回带层级前缀的命名参数

返回

名称形如 norm1.gamma、attention.query.weight、ffn1.weight 等

train()¶

递归切换到训练模式

注意

逐个调用 norm1、attention、norm2、两个 Linear 与 dropout 的 train()。

eval()¶

递归切换到推理模式

注意

推理时两处 dropout 均为恒等映射，注意力内部的 dropout 也一并关闭。

close()¶

释放全部子模块

注意

关闭后整个块不可再用；子模块由本块独占。

enum Activation.Kind¶

Activation 支持的激活函数种类

枚举常量共五个：RELU、GELU、SILU、SIGMOID、TANH，与 Tensor 上同名的一元算子一一对应。

枚举只作为 Activation 构造器的分派键，不含字段与方法；取值在 forward 中被 switch 穷举，因此新增常量需要同步修改分派逻辑。

注意

softmax 不在本枚举内；需要 softmax 时直接调用 Tensor.softmax(int)。

参见：basalt.torch.Activation, basalt.torch.Tensor#gelu()

枚举成员

RELU()¶

max(0, x)

GELU()¶

高斯误差线性单元

SILU()¶

Swish：x·sigmoid(x)

SIGMOID()¶

1 / (1 + e^(−x))

TANH()¶

双曲正切

LEAKY_RELU()¶

负半轴保留小斜率，避免神经元死亡

ELU()¶

负半轴指数衰减到 −α

static final class ComputationGraph.Builder¶

按名字登记节点、在 build 时排序并校验的计算图构造器

所有登记方法都返回 this，因此可以链式声明整张图。输入名与顶点名共用同一张命名空间表：addInputs 检查名字是否已作为输入或顶点存在，checkUnique 检查顶点名是否与已有输入或顶点冲突。节点登记不校验输入引用是否已声明，引用可以指向稍后才登记的节点，甚至形成环，这些检查全部推迟到 build。

build 先确认至少有一个输入和一个输出，再要求每个输出名指向输入或已登记顶点；接着按顶点输入个数建立入度表，从输入节点出发做 Kahn 拓扑排序得到执行顺序。排序结束后仍有顶点未出队，说明定义含环或存在从输入不可达的节点，两者都抛 IllegalArgumentException。构造过程不分配张量，也不执行任何计算。

示例

ComputationGraph graph = ComputationGraph.builder()
        .addInputs("x", "skip")
        .addAddVertex("res", "x", "skip")
        .addSubVertex("diff", "x", "skip")
        .setOutputs("res", "diff")
        .build();

注意

模块按引用登记，图与调用方共享同一个模块实例；ComputationGraph.close() 会关闭这些模块。

注意

顶点名与输入名不得重复，重复项在登记阶段即抛异常，不会留到 build。

注意

ADD 与 MUL 至少两个输入，SUB 固定两个；单输入模块节点用 addLayer 登记。

注意

同一个 Builder 可反复 build，每次得到独立实例，执行顺序在各自构建时确定。

参见：basalt.torch.ComputationGraph, basalt.torch.ComputationGraph.Vertex, basalt.torch.Module

方法

addInputs(String... names) → Builder¶

声明图输入节点，按声明顺序确定位置参数

逐个登记名字，并在写入前检查冲突，因此同一批参数内部的重复也会被发现。输入名同时决定 forwardMany 的位置顺序，也是 forward(Map) 的必需键。

参数

• names (String...) – 输入节点名，不得与已有输入或顶点重名

返回

this，便于链式调用

抛出

• IllegalArgumentException – 某个名字已作为输入或顶点登记

• NullPointerException – 任一名字为 null

addLayer(String name, Module<?, ?> module, String input) → Builder¶

登记引用模块的单输入节点

算子固定为 MODULE，顶点只保存模块引用，不复制实例；前向时取唯一输入调用 module.forward。模块参数按执行顺序进入整图的 parameters()，命名参数以“节点名.参数名”为键。

参数

• name (String) – 节点名，须在输入与顶点中唯一

• module (Module<?, ?>) – 被引用的模块，负责该节点的前向计算与参数

• input (String) – 上游节点或输入的名字，不要求已声明

返回

this，便于链式调用

抛出

• IllegalArgumentException – name 与已有输入或顶点重名

• NullPointerException – name、module 或 input 为 null

addAddVertex(String name, String... inputs) → Builder¶

登记逐元素相加节点，至少两个输入

算子为 ADD，不持有模块；前向时从第一个输入开始依次调用 add，因此参与相加的张量形状须能两两广播。残差连接就是把模块链的输入与输出同时登记为这里的两个输入。

参数

• name (String) – 节点名，须在输入与顶点中唯一

• inputs (String...) – 上游名字，至少两个，按传入顺序左结合相加

返回

this，便于链式调用

抛出

• IllegalArgumentException – inputs 为 null 或少于两个，或 name 与已有输入、顶点重名

• NullPointerException – inputs 中存在 null 元素

addSubVertex(String name, String leftInput, String rightInput) → Builder¶

登记左减右的二元减法节点

算子为 SUB，前向时计算 leftInput - rightInput，两侧形状须可广播。签名已限定两个输入，因此不再做输入个数检查。

参数

• name (String) – 节点名，须在输入与顶点中唯一

• leftInput (String) – 被减数的上游名字

• rightInput (String) – 减数的上游名字

返回

this，便于链式调用

抛出

• IllegalArgumentException – name 与已有输入或顶点重名

• NullPointerException – leftInput 或 rightInput 为 null

addMulVertex(String name, String... inputs) → Builder¶

登记逐元素相乘节点，至少两个输入

算子为 MUL，前向时从第一个输入开始依次调用 mul。常见用法是门控：把同一个上游名字分别作为门控信号与被门控特征传入，得到逐元素乘积。

参数

• name (String) – 节点名，须在输入与顶点中唯一

• inputs (String...) – 上游名字，至少两个，按传入顺序左结合相乘

返回

this，便于链式调用

抛出

• IllegalArgumentException – inputs 为 null 或少于两个，或 name 与已有输入、顶点重名

• NullPointerException – inputs 中存在 null 元素

setOutputs(String... names) → Builder¶

声明图输出节点，覆盖此前设置

先清空已有输出再逐个登记，重复调用以最后一次为准。输出名的存在性不在本方法校验，而是留到 build 时与输入名、顶点名一起比对。

注意

同名重复声明不会被去重，取回结果时按名字查表，因此得到的是同一个张量。

参数

• names (String...) – 输出名，须指向输入或已登记顶点

返回

this，便于链式调用

抛出

• NullPointerException – 任一名字为 null

build() → ComputationGraph¶

校验图定义并返回拓扑序固定的计算图

入度按每个顶点的输入个数初始化，再从输入节点出发逐个扣减消费方的入度，减到 0 即入队，出队顺序写入 executionOrder。校验与排序都只读构造器状态，因此构建成功后仍可继续登记节点，已建好的图不受影响。

返回

按拓扑序保存执行顺序的 ComputationGraph

抛出

• IllegalArgumentException – 没有输入或没有输出、输出名未知、顶点引用未知输入，或定义中存在环与不可达顶点

record ComputationGraph.Vertex¶

计算图节点的不可变描述：名字、算子、模块与输入名

四个分量在登记时确定：name 是节点标识，op 决定前向运算，module 只在 MODULE 节点上非 null，inputs 是上游名字的列表。op 为 MODULE 时 inputs 只有一个元素，ADD、SUB、MUL 节点的 module 为 null，只做张量运算。

inputs 由 List.of 生成，是不可变列表；record 的访问器直接返回该引用而不复制。节点按登记顺序存放在 Builder 的 LinkedHashMap 中，构图后经 vertices() 以只读映射暴露，遍历顺序与登记顺序一致，与 executionOrder() 的拓扑序不一定相同。

注意

record 只提供访问器，没有额外校验或归一化；取值前应先按 op 判断 module 是否可用。

注意

参数集合不包含在此记录中，参数由 module.parameters() 提供。

参见：basalt.torch.ComputationGraph, basalt.torch.ComputationGraph.VertexOp

enum ComputationGraph.VertexOp¶

节点算子类型：模块调用、相加、相减、相乘

MODULE 表示节点引用一个模块并对唯一输入调用 forward；ADD 与 MUL 把多个输入逐个累加或累乘，SUB 只接受两个输入计算左减右。前向的 switch 覆盖全部四个常量，没有默认分支。

枚举值同时决定顶点是否携带模块：只有 MODULE 节点的 module 非 null，其余三种节点的运算完全由张量算子完成，因此不贡献参数，也不随 train()、eval() 改变状态。

注意

ADD 与 MUL 的输入个数由 Builder 保证至少两个，枚举本身不携带元数信息。

注意

算子类型在登记后不可变，改变算子需要重新登记节点。

参见：basalt.torch.ComputationGraph.Vertex, basalt.torch.ComputationGraph.Builder

枚举成员

MODULE()¶

顶点执行一个 Module

ADD()¶

两个输入逐元素相加（残差）

SUB()¶

两个输入逐元素相减

MUL()¶

两个输入逐元素相乘（门控）