19 lines
508 B
Python
19 lines
508 B
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from collections import OrderedDict
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
def average_weights(weight_dicts: list[OrderedDict]) -> OrderedDict:
|
||
|
|
"""Average model state dicts from multiple workers."""
|
||
|
|
if not weight_dicts:
|
||
|
|
raise ValueError("No weights to average")
|
||
|
|
if len(weight_dicts) == 1:
|
||
|
|
return weight_dicts[0]
|
||
|
|
|
||
|
|
avg = OrderedDict()
|
||
|
|
for key in weight_dicts[0]:
|
||
|
|
avg[key] = torch.stack([w[key].float() for w in weight_dicts]).mean(dim=0)
|
||
|
|
return avg
|