7d11dfa688
Implementation of a mock scheduler for jobs/tasks that will evolve into a job manager for running training jobs over vast.ai instances.
18 lines
508 B
Python
18 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
|