Checkpoint Averaging

Problem

Average a list of weight checkpoints element-wise (a simple model soup): the mean over checkpoints for each parameter position.

The signature is checkpoint_average(checkpoints). Each checkpoint is a vector of the same length.

Hint

Average parameter by parameter across the checkpoint list.

case 1

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

expected: [2,3]

case 2

in: [[[0,0],[2,2],[4,4]]]

expected: [2,2]

case 3

in: [[[1]]]

expected: [1]

case 4

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

expected: [1,1]

Starter codePython

def checkpoint_average(checkpoints):
    # Your code here
    pass

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