THE BATCH-NORM TRAINING SWITCH

Difficulty: Medium Tags: Deep Learning, Batch Normalization, Training vs Inference

Story

A vision model serves predictions after training. A deployment bug keeps using batch statistics at inference, so the same image gets a different output depending on which images share its batch. Implement one-dimensional batch normalization with supplied running statistics and an explicit mode.

The Math

For a batch x_1,...,x_n, training mode uses the batch mean and population variance:

mu_B = (1/n) * sum_i(x_i)
v_B  = (1/n) * sum_i((x_i - mu_B)^2)

Inference mode uses the supplied running mean mu_R and running variance v_R. In either mode:

y_i = (x_i - mu) / sqrt(v + eps) * gamma + beta

The mode is 1 for training and 0 for inference.

Input Format

x (list of floats), mode (0 or 1), gamma, beta, eps, running_mean, running_var.

Output Format

List of normalized outputs (6 decimal places).

Constraints

• 1 <= n <= 10^5, eps > 0, running_var >= 0, magnitudes at most 10^6.
• Time limit: 1.0 second.

Example

Input

x = [1, 2, 3], mode = 0
gamma = 2, beta = 1, eps = 0.000001, running_mean = 2, running_var = 1

Output

[-0.999999, 1.000000, 2.999999]

Explanation: In inference mode, use running mean 2 and running variance 1, then apply scale 2 and bias 1.