BatchNorm Momentum Update

Problem

Update batch-norm running statistics: new = (1 - momentum) * old + momentum * batch, applied to both the mean and the variance.

The signature is batchnorm_momentum_update(running_mean, running_var, batch_mean, batch_var, momentum). Return [new_mean, new_var].

Hint

The same momentum coefficient is applied to the mean update and the variance update.

case 1

in: [0,1,1,2,0.1]

expected: [0.1,1.1]

case 2

in: [0,0,0,0,0.5]

expected: [0,0]

case 3

in: [2,4,3,9,0.25]

expected: [2.25,5.25]

case 4

in: [1,0.5,-1,0.1,0.9]

expected: [-0.8,0.14]

Starter codePython

def batchnorm_momentum_update(running_mean, running_var, batch_mean, batch_var, momentum):
    # Your code here
    pass

4 test cases · run your code against them in the editor