【RustyML入门】3.2. 全连接层与激活函数

【RustyML入门】3.2. 全连接层与激活函数

3.2. 全连接层与激活函数

Dense是前馈网络的主力层。它先做一次线性映射input * W + b,再套一个逐元素的非线性。RustyML 把这个非线性直接折叠进层本身,不会把它当成单独的一步来处理。

本页涉及构造函数及其折叠激活的设计,也涉及精确的初始化方案,以及五个激活函数各自的梯度与失效模式。还涉及由 GEMM 支撑的计算路径,以及如何读取或注入权重。本页假定你已经读过 Sequential 模型 那一页。这里讲的层,就是你往Sequential模型里堆叠的东西。

3.2.1. 折叠激活的设计

构造函数把激活函数作为第三个参数接收,而不是作为一个独立的层:

pub fn new( input_dim: usize, units: usize, activation: impl Into<Activation>, ) -> Result<Dense, Error>

层内部存的是一个Activation。这是一个只有 5 个变体、实现了Copy的普通枚举。它不是泛型参数Dense<A>,也不是Box<dyn Activation>。这是有意为之的选择。泛型参数会把Dense单态化成 5 份。

它还会逼着权重反序列化时逐个试探Dense<A>的每种组合,才能对上磁盘里的具体类型。trait 对象则会给每次逐元素调用都加一层间接跳转。运行时枚举让Dense始终保持为单一的具体类型。于是持久化层能把每个存下来的层向下转型到唯一一个结构体(见权重保存与加载)。激活的数学运算则始终是一个纯粹、无状态的函数,由层在自己的前向和反向传播里调用它。

Keras 用Dense(units, activation="relu")做了同样的折叠选择。RustyML 没有隐式的“不带激活”。你必须始终显式传一个。想要一个纯线性层,就传Activation::Linear,也就是恒等映射。这对应的正是 Keras 的activation=None。下面每个以回归输出头收尾的示例都是这么做的。

折叠是合理的默认选择。RustyML 也把这些激活函数提供为独立的层(ReLUSigmoidTanhSoftmaxLinear)。当你需要在线性映射和非线性之间插入别的东西时,就用这些独立层。最常见的情况就是插一个归一化层(见正则化与归一化层)。两种写法是等价的:

// 这两种堆叠方式计算结果相同。 model.add(Dense::new(64, 32, Activation::ReLU).unwrap()); // 折叠写法:只缓存一份激活输出 model.add(Dense::new(64, 32, Activation::Linear).unwrap()) // 拆分写法:先来一个恒等 Dense…… .add(ReLU::new()); // ……再接一个独立的 ReLU

优先用折叠写法。它只缓存一份激活后的张量,而不是两份。反向传播随后就基于这份缓存输出对激活求导(见 3.2.5)。只有当某个层必须待在W*x + b和非线性之间时,才用拆分写法。

3.2.2. 构造并确定层的尺寸

Dense::new(input_dim, units, activation)返回Result<Dense, Error>。两个维度都必须非零。任意一个传0都会得到Error::InvalidParameter(见错误处理)。input_dim是每行的特征数。units是神经元个数,也就是输出的宽度。

参数量是input_dim * units + units:每个(输入,输出)对一个权重,再加每个输出一个偏置。param_count()把它报告为TrainingParameters::Trainable(n)output_shape()则渲染成(None, units),这里的None是动态的 batch 维度,跟 Keras 的 summary 一致。所以一个Dense(4, 3, ...)持有一个4 x 3的权重矩阵和一个1 x 3的偏置,共12 + 3 = 15个可训练标量。

输入必须是形状为(batch, input_dim)的二维张量。一维或三维输入会得到Error::InvalidInput,而不会被悄悄 reshape。如果你喂进来的是卷积或循环层堆叠的输出,先把它展平成两维。下面这段程序构造一个层,检视它,再注入一组已知权重:

usendarray::Array2;userustyml::neural_network::layers::layer_weight::LayerWeight;userustyml::neural_network::layers::{Activation,Dense,TrainingParameters};userustyml::neural_network::traits::Layer;fnmain(){// 一个 4 -> 3 的全连接层,折叠了 ReLU 激活。letmutdense=Dense::new(4,3,Activation::ReLU).unwrap();// 4*3 个权重 + 3 个偏置 = 15 个可训练标量。assert_eq!(dense.param_count(),TrainingParameters::Trainable(15));println!("output shape: {}",dense.output_shape());// (None, 3)// 读取刚初始化好的参数,且不发生克隆。matchdense.get_weights(){LayerWeight::Dense(w)=>{println!("weight {:?}, bias {:?}",w.weight.shape(),w.bias.shape());}_=>unreachable!(),}// 注入已知权重;形状会对照层的配置做校验。letweights=Array2::from_shape_vec((4,3),(0..12).map(|v|vasf32).collect::<Vec<f32>>()).unwrap();letbias=Array2::zeros((1,3));dense.set_weights(weights,bias).unwrap();}

set_weights会检查两个形状,一旦不匹配就返回Error::NeuralNetwork(NnError::WeightShape)。给一个4 -> 3的层传(3, 3)的权重,或者传(1, 4)的偏置,都会被拒绝,而不是被截断。

3.2.3. 权重初始化

权重采用Xavier/Glorot 均匀初始化。每个元素从Uniform(-limit, +limit)中抽取,其中limit = sqrt(6 / (input_dim + units))。偏置从恰好为零开始。这就是把 fan-in 和 fan-out 相加的 Glorot 方案。它跟 Keras 的Dense默认设置完全一致(glorot_uniform权重、zeros偏置)。

RustyML 无论激活函数是什么都套用 Glorot。它不会为ReLU层切换到 He/Kaiming 初始化,尽管 He 才是教科书里跟整流器配套的选择。对浅层网络这基本无所谓。对很深的 ReLU 堆叠,早期收敛可能会比 He 初始化的等价网络略慢。真遇到这种情况,就用set_weights传入你自己缩放好的抽样来初始化。

初始化是通过 crate 共享的 RNG 抽样的。默认情况下,如果设了进程级全局种子就用它,否则就取自系统熵。这样一来,除非你固定住随机性,两次运行会得到不同的权重。想让它可复现有 2 种办法。一种是在构造模型之前,用rustyml::random::set_global_seed(...)设一个线程局部的全局种子。这同时也固定了 dropout 掩码和 fit 时的 batch 打乱。

另一种办法是用Dense::new(...)?.with_random_state(seed)给单个层显式播种。它会用那个种子重跑一遍 Glorot,而不动全局的随机流。完整的播种模型见可复现性与随机种子。那里也解释了为什么给某一层显式播种不会改变分发给未播种层的种子。

3.2.4. 五个激活函数

Activation恰好有 5 个变体。这个 crate 里没有 LeakyReLU、ELU、GELU 或 Swish。每个变体的反向传播都是用激活后的输出a = f(x)来表达的,而不是激活前的x,因为层缓存的正是前者。

Activation前向f(x)反向(给定上游g、输出a输出范围
ReLUmax(0, x)a > 0处透传g,否则为0[0, inf)
Sigmoid1 / (1 + e^-x)g * a * (1 - a)(0, 1)
Tanhtanh(x)g * (1 - a^2)(-1, 1)
Softmax平移后的exp,按行归一化a_i * (g_i - sum_j(a_j * g_j))(按行的 Jacobian)单纯形,每行和为1
Linearx原样透传g(-inf, inf)

ReLU是隐藏层的默认选择。它便宜,而且在正半轴不饱和。它的失效模式是死神经元。因为x <= 0处导数为0,一个在整个 batch 的每个样本上激活前值都为负的神经元,拿到的梯度就是零。

它永远不会更新,从此彻底关死。学习率过高会让情况更糟。它会早早把神经元推进死区。Glorot 初始化没有 leaky 变体可用。你能用的办法是更小的学习率和缩放得当的输入,或者,如果是你自己注入权重,用更合理的初始缩放。

Sigmoid把值压到(0, 1)。它的导数a * (1 - a)a = 0.5时取到峰值0.25,并随着输出饱和而衰减到零。在很深的隐藏路径上叠 sigmoid 会这样扼住梯度:这就是经典的梯度消失问题。把Sigmoid用在单个二分类输出(搭配二分类交叉熵)或者当作门控。别拿它当很深的隐藏非线性。输入极端时,f32会把它饱和到恰好的0.01.0

Tanh映射到(-1, 1)。跟 sigmoid 不同,它以零为中心,这往往能让隐藏层的优化更规矩。它的梯度1 - a^2在原点附近仍能取到1。在两端它跟 sigmoid 一样会饱和。它是那种天然有界、以零为中心的隐藏激活,循环层内部用的也是它。

Softmax把一行 logits 变成最后一个轴上的概率分布。前向传播在取指数之前,先减去每行的最大值。这让最大的那一项变成exp(0) = 1,总和也就始终>= 1。这个平移让 softmax 既防溢出又平移不变:给每个 logit 加同一个常数,输出不变。

Softmax 至少需要二维输入。一维张量会返回Error::InvalidInput。它的反向传播是整行上真正的 Jacobian 向量积,而不是逐元素相乘,每行梯度之和为零。

Softmax 属于分类器的输出层,要搭配正确的损失函数配置。夹在网络中间时,Softmax 在机理上是合法的:Jacobian 反向是正确的,所以埋在堆叠里的Dense(..., Softmax)照样能训练。但这几乎从来不是你想要的。它把表示塌缩到一个单纯形上,丢掉所有幅度信息,而且比 ReLU 或 tanh 饱和得更厉害。

Linear就是恒等映射。它的梯度为1,输出无界。把它用在回归输出头上,以及任何损失函数期望原始 logits 的场合。

损失函数必须匹配输出头上的激活函数。多分类有 2 种正确的搭配。把它们混起来会悄无声息地毁掉训练:

  • Softmax 输出头 +CategoricalCrossEntropy::new(false)最后的Dense吐出概率。损失函数吃进概率。正确。
  • Linear 输出头 +CategoricalCrossEntropy::new(true)最后的Dense吐出原始 logits。损失函数在内部做一次数值稳定的 log-softmax,一步返回折叠好的(softmax(z) - y)梯度。同样正确,而且数值上更稳。

这 2 种搭配在数学上产生相同的梯度。折叠的from_logits = true路径避开了中间的-y/p除法和单独的一次 softmax 步骤。所以当某个预测概率极小时,它劣化得更平缓。不要把这两种搭配混在一起。

softmax 输出头配from_logits = true会做两次 softmax。linear 输出头配from_logits = false会把原始 logits 喂给一个期望概率的损失函数。损失函数的细节见损失函数。

// Softmax 输出头之外更稳定的替代方案:输出 logits,把 softmax 折叠进损失函数。 model .add(Dense::new(8, 3, Activation::Linear).unwrap()) // 原始 logits,不是概率 .compile(optimizer, CategoricalCrossEntropy::new(true)); // from_logits = true

这些激活函数是纯数学,不做任何NaN/Inf清洗。NaN会原样传下去。tanh在输入很大时会饱和到1-1ReLU对很大的负输入结果是0。softmax 一行里任何位置出现的NaN,都会通过归一化因子污染整行。非有限值是在下游浮现的,表现为NaN损失,而不是在激活这一步。

3.2.5. 前向、反向与 GEMM 路径

前向传播是activation(input * W + b)。它是作为对 gemmkit 后端的一次调用算出来的。乘积、按列的偏置,以及(对ReLU而言)激活函数,全都在一趟里跑完。偏置和激活是在输出分块还留在寄存器里的时候,由内核的尾声(epilogue)施加的。

融合后的结果与分开做乘积再加标量激活相比,逐位相同,只有一个例外。融合的ReLU尾声会把NaN的激活前值映射成0.0。crate 里其他地方用的独立ReLU激活则会传播NaNNaN的激活前值本身就意味着模型已经发散。缓存输出版本的ReLU导数接下来会把那个0.0当成死神经元处理。它在那里反传回一个零梯度,而不是继续把NaN往回传。

乘积是否开线程由 gemmkit 决定,而不是这一层决定。它拿batch * input_dim * units去和一个工作量闸门比,默认是 589,824(这是原始乘积本身,不是带了 2 倍系数的 FLOP 数)。低于这个闸门,乘积就留在单线程上。高于它,worker 数会随工作量逐步爬升,而不是一步跳到整台机器的宽度。

紧凑循环里的小层按设计保持串行,因为每次调用的派发开销会占大头。这些乘积嵌在外层并行区域里也能安全嵌套。后端在矩阵乘法里有完整描述,可调的阈值则在性能调优与并行。

紧随其后的逐元素激活有它自己的、独立的并行阈值。ReLU是一个受内存带宽限制的“廉价映射”,它的临界点是 4,000,000 个元素。在任何实际的层规模下,它都串行运行。SigmoidTanhSoftmaxexp主导,超过 131,072 个元素就转并行。挪动这些阈值只是在串行和并行之间做交换。结果完全相同,而且在同一台机器上,每次运行每个乘积都是确定一致的。

forward会为反向传播缓存输入和激活后的输出。predict是它的求值模式孪生。它算出同样的值,但不写任何缓存。反向传播先用缓存的输出对激活求导。

接下来它算 3 个量。权重梯度是input^T * grad,一次 GEMM。偏置梯度是按 batch 做的列求和。输入梯度是grad * W^T,又一次 GEMM。

forward之前调用backward会返回Error::NeuralNetwork(NnError::ForwardPassNotRun)。形状对不上缓存输出的上游梯度会返回Error::ShapeMismatch。两者都是错误,绝不会 panic。

3.2.6. 读取与设置权重

没有单独的weights()/bias()getter。访问器是Layertrait 里的get_weights()。它返回一个LayerWeight枚举。对全连接层来说,那就是LayerWeight::Dense(DenseLayerWeight { weight, bias }),其中weightbias是从活着的层借来的Cow<Array2<f32>>,不发生克隆。

weight的形状是(input_dim, units)bias的形状是(1, units),正如 3.2.2 里的构造示例所示。要写入参数,用set_weights(weights, bias),它会校验两个形状。这同一个LayerWeight枚举就是磁盘上的权重格式。你在这里能读到的东西,正是权重保存与加载里存取时来回往返的东西。

3.2.7. 两个实战模型

回归网络以一个Linear输出头收尾,针对均方误差训练。隐藏层折叠ReLU。输出层折叠Linear,因为回归目标是无界的:

usendarray::Array;userustyml::neural_network::layers::{Activation,Dense};userustyml::neural_network::losses::MeanSquaredError;userustyml::neural_network::optimizers::SGD;userustyml::neural_network::sequential::Sequential;fnmain(){// 4 个样本,3 个特征,每个样本 1 个连续目标值。letx=Array::from_shape_vec((4,3),vec![0.0,0.1,0.2,1.0,0.9,0.8,0.2,0.1,0.0,0.9,1.0,0.8],).unwrap().into_dyn();lety=Array::from_shape_vec((4,1),vec![0.3,2.7,0.3,2.7]).unwrap().into_dyn();letmutmodel=Sequential::new();model.add(Dense::new(3,8,Activation::ReLU).unwrap())// 隐藏层,折叠 ReLU.add(Dense::new(8,1,Activation::Linear).unwrap())// 回归输出头:恒等.compile(SGD::new(0.05,0.9,false,0.0).unwrap(),MeanSquaredError::new());model.fit(&x,&y,20).unwrap();letpreds=model.predict(&x).unwrap();println!("prediction shape: {:?}",preds.shape());// [4, 1]}

分类器以一个作用在 one-hot 目标上的Softmax输出头收尾,搭配CategoricalCrossEntropy::new(false),因为这个输出头吐出的是概率:

usendarray::Array;userustyml::neural_network::layers::{Activation,Dense};userustyml::neural_network::losses::CategoricalCrossEntropy;userustyml::neural_network::optimizers::Adam;userustyml::neural_network::sequential::Sequential;fnmain(){// 4 个样本,3 个特征,2 个类别(one-hot 目标)。letx=Array::from_shape_vec((4,3),vec![0.0,0.1,0.2,1.0,0.9,0.8,0.1,0.0,0.2,0.8,1.0,0.9],).unwrap().into_dyn();lety=Array::from_shape_vec((4,2),vec![1.0,0.0,0.0,1.0,1.0,0.0,0.0,1.0]).unwrap().into_dyn();letmutmodel=Sequential::new();model.add(Dense::new(3,8,Activation::ReLU).unwrap()).add(Dense::new(8,2,Activation::Softmax).unwrap())// 概率输出头.compile(Adam::new(0.01,0.9,0.999,1e-8,0.0).unwrap(),CategoricalCrossEntropy::new(false),);model.fit(&x,&y,20).unwrap();letprobs=model.predict(&x).unwrap();// 每一行都是 2 个类别上的分布,和为 1。println!("class-probability shape: {:?}",probs.shape());// [4, 2]}

把输出头换成Activation::Linear,损失换成CategoricalCrossEntropy::new(true)。这个分类器就会走上更稳定的折叠-logits 路径,predict产出的是原始分数而不是概率。该用哪个优化器来 compile,以及学习率和动量如何跟上面的死神经元、饱和行为相互作用,是优化器的主题。