SWA Model Average

Problem

Stochastic weight averaging takes the elementwise mean of a list of parameter vectors.

Given the list of models, return the averaged model.

Hint

Every checkpoint contributes equally.

case 1

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

expected: [2,3]

case 2

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

expected: [1,1]

case 3

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

expected: [3]

Starter codePython

def swa_model_average(models):
    # Your code here
    pass

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