From 051d381be5e23924ee047c30d47cdad085ff16a3 Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Thu, 29 Oct 2026 17:16:58 -0400 Subject: [PATCH 1/3] fix(densenet): stop holding a forward's activations after it returns - DenseBlockLayer and TransitionLayer wrote every intermediate (BN, ReLU, conv outputs) to fields that nothing reads; the tape keeps what backward needs. They are locals now. - DenseBlock kept a list of every inner layer's output that nothing reads; removed. - BatchNormalizationLayer cached its input whenever no tape was recording, so in eval mode it held every BN input of the last forward. It now also requires ShouldCacheForBackward. - PredictAndValidateChain's layer observer records each layer's input and output; it now drops that recording after validating it, so a thread or task that captured the observer during the forward cannot keep those tensors alive (ooples/AiDotNet.Tensors#1610 fixes the capture itself). DenseNet-121, CPU, b16 inference, WS: 1904 -> 1617 MB (with Tensors #1610). Part of #2383. Co-Authored-By: Claude Opus 5.5 --- .../Graph/LayerForwardObserver.cs | 5 +++- .../Layers/BatchNormalizationLayer.cs | 6 ++-- src/NeuralNetworks/Layers/DenseBlock.cs | 7 +---- src/NeuralNetworks/Layers/DenseBlockLayer.cs | 28 ++++++------------- src/NeuralNetworks/Layers/TransitionLayer.cs | 28 ++++++------------- src/NeuralNetworks/NeuralNetworkBase.cs | 7 +++++ 6 files changed, 33 insertions(+), 48 deletions(-) diff --git a/src/NeuralNetworks/Graph/LayerForwardObserver.cs b/src/NeuralNetworks/Graph/LayerForwardObserver.cs index 6403cdf83e..3c6af42f9a 100644 --- a/src/NeuralNetworks/Graph/LayerForwardObserver.cs +++ b/src/NeuralNetworks/Graph/LayerForwardObserver.cs @@ -1,4 +1,4 @@ -using System; +using System; using System.Collections.Generic; using System.Runtime.CompilerServices; using System.Threading; @@ -216,6 +216,9 @@ public bool HasFanOut } } + /// Drops the recorded calls and the tensors they reference, once they have been read. + public void ReleaseRecording() => _calls.Clear(); + /// public void Dispose() { diff --git a/src/NeuralNetworks/Layers/BatchNormalizationLayer.cs b/src/NeuralNetworks/Layers/BatchNormalizationLayer.cs index 718dcdd84d..99298e449e 100644 --- a/src/NeuralNetworks/Layers/BatchNormalizationLayer.cs +++ b/src/NeuralNetworks/Layers/BatchNormalizationLayer.cs @@ -869,10 +869,10 @@ protected override Tensor ForwardTraced(Tensor input) // that's never reached when training goes through the tape (tape // holds its own intermediate refs already). Skip the assignment // when a tape is active so this field doesn't double-root the - // input activation; null it out so the previous step's tensor is - // eligible for collection. + // input activation, and in inference (ShouldCacheForBackward is false + // in eval mode), where it held every BN input of the last forward. bool tapeActive = AiDotNet.Autodiff.GradientRecording.IsRecording(); - _lastInput = tapeActive ? null : input; + _lastInput = tapeActive || !ShouldCacheForBackward ? null : input; // Batch statistics are taken over every value a feature has in this batch: the batch // axis AND, for convolutional (channels-first) inputs, every spatial position. That is diff --git a/src/NeuralNetworks/Layers/DenseBlock.cs b/src/NeuralNetworks/Layers/DenseBlock.cs index 7f70c94241..60d3baff80 100644 --- a/src/NeuralNetworks/Layers/DenseBlock.cs +++ b/src/NeuralNetworks/Layers/DenseBlock.cs @@ -1,4 +1,4 @@ -using AiDotNet.ActivationFunctions; +using AiDotNet.ActivationFunctions; using AiDotNet.Attributes; using AiDotNet.Autodiff; using AiDotNet.Interfaces; @@ -123,7 +123,6 @@ public partial class DenseBlock : LayerBase, ILayerSerializationExtras, // Non-readonly: lazy ctor leaves _inputChannels = -1 until // OnFirstForward resolves it from the runtime input tensor. private int _inputChannels; - private List>? _layerOutputs; public override bool SupportsTraining => true; @@ -261,8 +260,6 @@ protected override Tensor ForwardTraced(Tensor input) // before the inner-layer Forwards consume their weights. if (!IsShapeResolved) OnFirstForward(input); - _layerOutputs = new List>(_numLayers + 1) { input }; - // Current feature maps (accumulated) var currentFeatures = input; @@ -270,7 +267,6 @@ protected override Tensor ForwardTraced(Tensor input) { // Each layer takes ALL previous features as input var layerOutput = layer.Forward(currentFeatures); - _layerOutputs.Add(layerOutput); // Concatenate new features with existing features along channel dimension currentFeatures = ConcatenateChannels(currentFeatures, layerOutput); @@ -322,7 +318,6 @@ internal override Dictionary GetMetadata() /// public override void ResetState() { - _layerOutputs = null; foreach (var layer in _layers) { layer.ResetState(); diff --git a/src/NeuralNetworks/Layers/DenseBlockLayer.cs b/src/NeuralNetworks/Layers/DenseBlockLayer.cs index 4e4d07ae1a..4446721d3a 100644 --- a/src/NeuralNetworks/Layers/DenseBlockLayer.cs +++ b/src/NeuralNetworks/Layers/DenseBlockLayer.cs @@ -1,4 +1,4 @@ -using AiDotNet.ActivationFunctions; +using AiDotNet.ActivationFunctions; using AiDotNet.Attributes; using AiDotNet.Autodiff; using AiDotNet.Interfaces; @@ -74,13 +74,6 @@ public partial class DenseBlockLayer : LayerBase, ILayerSerializationExtra [Scratch] private Tensor? _lastInput; - private Tensor? _bn1Out; - private Tensor? _relu1Out; - private Tensor? _conv1Out; - [AiDotNet.Attributes.Scratch] - private Tensor? _bn2Out; - [AiDotNet.Attributes.Scratch] - private Tensor? _relu2Out; public override bool SupportsTraining => true; @@ -203,17 +196,19 @@ protected override Tensor ForwardTraced(Tensor input) var x = input.Shape.Length == 3 ? Engine.Reshape(input, [1, input.Shape[0], input.Shape[1], input.Shape[2]]) : input; // BN-ReLU-Conv1x1 (DenseNet paper: pre-activation bottleneck) - _bn1Out = _bn1.Forward(x); + // The intermediates are locals: the tape keeps what backward needs, and a field would hold + // every dense layer's activations after the forward, in inference too (DenseNet-121 b16). + var bn1Out = _bn1.Forward(x); // Scalar activation objects operate eagerly and therefore detach the value from the // active GradientTape. Keep the DenseNet pre-activation on the engine graph so gradients // reach BN1 and every layer before this block. - _relu1Out = Engine.ReLU(_bn1Out); - _conv1Out = _conv1x1.Forward(_relu1Out); + var relu1Out = Engine.ReLU(bn1Out); + var conv1Out = _conv1x1.Forward(relu1Out); // BN-ReLU-Conv3x3 - _bn2Out = _bn2.Forward(_conv1Out); - _relu2Out = Engine.ReLU(_bn2Out); - var output = _conv3x3.Forward(_relu2Out); + var bn2Out = _bn2.Forward(conv1Out); + var relu2Out = Engine.ReLU(bn2Out); + var output = _conv3x3.Forward(relu2Out); // Remove batch dim if we added it if (input.Shape.Length == 3 && output.Shape.Length == 4 && output.Shape[0] == 1) @@ -269,11 +264,6 @@ private void ApplyParameters(Vector parameters) public override void ResetState() { _lastInput = null; - _bn1Out = null; - _relu1Out = null; - _conv1Out = null; - _bn2Out = null; - _relu2Out = null; _bn1.ResetState(); _conv1x1.ResetState(); diff --git a/src/NeuralNetworks/Layers/TransitionLayer.cs b/src/NeuralNetworks/Layers/TransitionLayer.cs index 28bbb029bb..66ce266f33 100644 --- a/src/NeuralNetworks/Layers/TransitionLayer.cs +++ b/src/NeuralNetworks/Layers/TransitionLayer.cs @@ -1,4 +1,4 @@ -using AiDotNet.ActivationFunctions; +using AiDotNet.ActivationFunctions; using AiDotNet.Attributes; using AiDotNet.Autodiff; using AiDotNet.Interfaces; @@ -82,11 +82,6 @@ public partial class TransitionLayer : LayerBase, ILayerSerializationExtra private readonly AveragePoolingLayer _pool; private readonly IActivationFunction _relu; - [Scratch] - private Tensor? _lastInput; - private Tensor? _bnOut; - private Tensor? _reluOut; - private Tensor? _convOut; /// /// Stores the original input shape for any-rank tensor support. @@ -304,19 +299,18 @@ protected override Tensor ForwardTraced(Tensor input) throw new ArgumentException($"TransitionLayer requires at least 3D input, got {rank}D"); } - _lastInput = processInput; - - // BN → ReLU → Conv1x1 → AvgPool - _bnOut = _bn.Forward(processInput); - _reluOut = _relu.Activate(_bnOut); - _convOut = conv.Forward(_reluOut); + // BN → ReLU → Conv1x1 → AvgPool. Locals, not fields: nothing reads them after the forward, + // and a field would keep this layer's activations alive between forwards. + var bnOut = _bn.Forward(processInput); + var reluOut = _relu.Activate(bnOut); + var convOut = conv.Forward(reluOut); // Handle batched input (4D) - use Engine.AvgPool2D directly // AveragePoolingLayer expects 3D input, so we handle 4D separately // [B, C, H, W] uses Engine directly, [C, H, W] uses the AveragePoolingLayer - Tensor output = _convOut.Shape.Length == 4 - ? Engine.AvgPool2D(_convOut, poolSize: 2, stride: 2, padding: 0) - : _pool.Forward(_convOut); + Tensor output = convOut.Shape.Length == 4 + ? Engine.AvgPool2D(convOut, poolSize: 2, stride: 2, padding: 0) + : _pool.Forward(convOut); // Restore original batch dimensions for any-rank support if (_originalInputShape != null && _originalInputShape.Length > 4) @@ -429,10 +423,6 @@ internal override Dictionary GetMetadata() /// public override void ResetState() { - _lastInput = null; - _bnOut = null; - _reluOut = null; - _convOut = null; _bn.ResetState(); _conv?.ResetState(); diff --git a/src/NeuralNetworks/NeuralNetworkBase.cs b/src/NeuralNetworks/NeuralNetworkBase.cs index 6120af0c39..2a008ad417 100644 --- a/src/NeuralNetworks/NeuralNetworkBase.cs +++ b/src/NeuralNetworks/NeuralNetworkBase.cs @@ -6324,6 +6324,13 @@ private Tensor PredictAndValidateChain(Tensor input) { // Diagnostic only. } + finally + { + // The recording holds every layer's input and output of that forward. Drop it here, so a + // context that captured the observer (a thread or task started during the forward) cannot + // keep those activations alive. + trace.ReleaseRecording(); + } } return output; From ebda8208c0af8ad0ae53526774602a6901f25ffd Mon Sep 17 00:00:00 2001 From: Franklin Moormann Date: Thu, 29 Oct 2026 20:44:46 -0400 Subject: [PATCH 2/3] style(densenet): restore utf-8 bom on denseblock sources Co-Authored-By: Claude Opus 5.5 --- src/NeuralNetworks/Layers/DenseBlock.cs | 2 +- src/NeuralNetworks/Layers/DenseBlockLayer.cs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/NeuralNetworks/Layers/DenseBlock.cs b/src/NeuralNetworks/Layers/DenseBlock.cs index 60d3baff80..fb90c7e6e8 100644 --- a/src/NeuralNetworks/Layers/DenseBlock.cs +++ b/src/NeuralNetworks/Layers/DenseBlock.cs @@ -1,4 +1,4 @@ -using AiDotNet.ActivationFunctions; +using AiDotNet.ActivationFunctions; using AiDotNet.Attributes; using AiDotNet.Autodiff; using AiDotNet.Interfaces; diff --git a/src/NeuralNetworks/Layers/DenseBlockLayer.cs b/src/NeuralNetworks/Layers/DenseBlockLayer.cs index 4446721d3a..35bd824da3 100644 --- a/src/NeuralNetworks/Layers/DenseBlockLayer.cs +++ b/src/NeuralNetworks/Layers/DenseBlockLayer.cs @@ -1,4 +1,4 @@ -using AiDotNet.ActivationFunctions; +using AiDotNet.ActivationFunctions; using AiDotNet.Attributes; using AiDotNet.Autodiff; using AiDotNet.Interfaces; From a53ba5326b484a829fb57bf1bc26a0eab975e750 Mon Sep 17 00:00:00 2001 From: t Date: Fri, 30 Oct 2026 08:57:46 -0400 Subject: [PATCH 3/3] fix(nn): release the trace recording when the observed forward or the layer check throws Co-Authored-By: Claude Opus 5.5 --- src/NeuralNetworks/NeuralNetworkBase.cs | 34 +++++++++++++++++-------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/src/NeuralNetworks/NeuralNetworkBase.cs b/src/NeuralNetworks/NeuralNetworkBase.cs index 2a008ad417..93a493082d 100644 --- a/src/NeuralNetworks/NeuralNetworkBase.cs +++ b/src/NeuralNetworks/NeuralNetworkBase.cs @@ -6293,6 +6293,12 @@ private Tensor PredictAndValidateChain(Tensor input) { observed = InferenceArenaSettings.Enabled ? Predict(input) : PredictCore(input); } + catch + { + // The recording holds every layer's activations; a failed forward must not leave them alive. + trace?.ReleaseRecording(); + throw; + } finally { trace?.Dispose(); @@ -6303,16 +6309,24 @@ private Tensor PredictAndValidateChain(Tensor input) // untraced path, and a second forward would draw a stochastic layer's randomness twice (a // seeded Monte Carlo dropout call then differed from the next one). Under an enclosing // observer the re-run would take the same unfolded path, so it is skipped there too. - output = trace is null || !trace.SuppressedFolding || enclosedByObserver - ? observed - : (InferenceArenaSettings.Enabled ? Predict(input) : PredictCore(input)); - - // AFTER the forward, deliberately. Measured with a probe rather than assumed: at every - // earlier point tried - ReportLayerContractMismatches, ResolveLazyLayerShapes, and the top - // of this method's caller - Layers was still EMPTY, because these stacks materialise during - // the forward. A check that runs before the layers exist cannot fail, and would have shipped - // as a diagnostic that silently never fired. - ReportArchitectureLayerDisagreement(); + try + { + output = trace is null || !trace.SuppressedFolding || enclosedByObserver + ? observed + : (InferenceArenaSettings.Enabled ? Predict(input) : PredictCore(input)); + + // AFTER the forward, deliberately. Measured with a probe rather than assumed: at every + // earlier point tried - ReportLayerContractMismatches, ResolveLazyLayerShapes, and the top + // of this method's caller - Layers was still EMPTY, because these stacks materialise during + // the forward. A check that runs before the layers exist cannot fail, and would have shipped + // as a diagnostic that silently never fired. + ReportArchitectureLayerDisagreement(); + } + catch + { + trace?.ReleaseRecording(); + throw; + } if (trace is not null) {