paddle.sparse.nn.BatchNorm keeps the mode of its first forward: eval after training uses batch statistics and updates the stored ones #79945

bug描述 Describe the Bug

paddle.sparse.nn.BatchNorm created with the default use_global_stats=None should decide on every call: batch statistics in train mode, stored statistics in eval mode (its docstring: "if set to None, use global statistics in the test phase and use the statistics of one mini-batch in the training phase"). Instead the first forward writes its decision onto the layer, and every later forward keeps it:

• first forward in train mode, then eval() (the usual order): eval normalises with the statistics of the batch and overwrites the stored mean / variance;
• first forward in eval mode, then train(): training normalises with the stored statistics and never updates them.

No warning or error is raised. paddle.nn.BatchNorm1D fed the same values behaves correctly in both orders.

Minimal reproduction:

import numpy as np
import paddle

rng = np.random.RandomState(0)


def batch(shift):  # a sparse input with no zero entry
    x = (rng.randn(2, 4, 4, 4, 3) + shift).astype("float32")
    return paddle.to_tensor(x).to_sparse_coo(4)


# 1) the usual order: train, then eval
bn = paddle.sparse.nn.BatchNorm(3)  # use_global_stats=None (the default)
bn.train()
bn(batch(1.0))
print("after a train forward: _use_global_stats =", bn._use_global_stats)
bn.eval()
before = bn._mean.numpy().copy()
bn(batch(-2.0))
print("eval forward moved the stored mean by", float(np.abs(bn._mean.numpy() - before).max()))

# 2) eval first, then train
bn = paddle.sparse.nn.BatchNorm(3)
bn.eval()
bn(batch(1.0))
print("after an eval forward: _use_global_stats =", bn._use_global_stats)
bn.train()
before = bn._mean.numpy().copy()
bn(batch(1.0))
print("train forward moved the stored mean by", float(np.abs(bn._mean.numpy() - before).max()))

Output on 3.3.1 (same on GPU and CPU):

after a train forward: _use_global_stats = False
eval forward moved the stored mean by 0.22993075847625732
after an eval forward: _use_global_stats = True
train forward moved the stored mean by 0.0

Expected: 0.0 for the eval forward and a non-zero value for the train forward, with _use_global_stats still None.

A fuller check against paddle.nn.BatchNorm1D on the same values (train, train, eval): the sparse layer's eval output matches the batch-statistics formula exactly and differs from the stored-statistics formula by up to 3.93; the same sample placed in two different batches gives eval outputs that differ by up to 2.06. The dense layer gives 0.0 for both.

Where: python/paddle/sparse/nn/layer/norm.py, BatchNorm.forward:

if self._use_global_stats is None:
    self._use_global_stats = not self.training
    trainable_statistics = False
else:
    trainable_statistics = not self._use_global_stats

After the first call the attribute is no longer None, so later calls take the else branch: after a train-first call use_global_stats=False, trainable_statistics=True is passed even in eval mode; after an eval-first call use_global_stats=True is passed even in train mode. paddle.nn.functional.batch_norm has the same lines on a local variable, which is why the dense layers are not affected.

Impact: evaluation results of a sparse model depend on the composition of the eval batch, and evaluating changes the stored statistics. A model that is evaluated once before training silently trains with frozen BatchNorm statistics. Code that reads _use_global_stats to see how the layer was configured gets the mode of the first call instead (this came up in #79944, where the conversion could not tell a sparse layer created with use_global_stats=True from one that had only run in eval mode).

Proposed fix: keep the decision in a local variable and pass that to the op (3 places in forward), as paddle.nn.functional.batch_norm does. I tried this change on the develop file: the sparse layer then matches paddle.nn.BatchNorm1D in both orders, and the 7 existing tests in test_sparse_norm_op.py still pass. I can open a PR with tests. Note that it changes the eval results of existing sparse models (they start using the stored statistics), so the PR would be marked as changing numeric results.

其他补充信息 Additional Supplementary Information

• PaddlePaddle 3.3.1 (GPU wheel built for CUDA 12.6 / cuDNN 9.9), Python 3.12, Windows 11; reproduced on GPU and on CPU.
• The lines are the same on develop (checked at commit 77bc015); reproduced with that file loaded in 3.3.1.
• Not checked: static graph mode (the same attribute is passed there), and paddle.sparse.nn.SyncBatchNorm, which does not use this attribute.