Gradient Accumulation Mean

Problem

Gradient accumulation sums micro-batch gradients and divides by the number of accumulation steps to match a large batch.

Given a list of micro-batch gradients and the number of accumulation steps, return the averaged gradient computed as the sum divided by the step count.

Hint

Divide the accumulated sum by the configured step count, not the list length.

Test cases

case 1

in: [[[1,2],[3,4]],2]

expected: [2,3]

case 2

in: [[[2],[4],[6]],3]

expected: [4]

case 3

in: [[[1,1],[1,1]],4]

expected: [0.5,0.5]

Starter codePython

def gradient_accumulation_mean(grads, steps):
    # Your code here
    pass

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