58 lines
1.3 KiB
Python
58 lines
1.3 KiB
Python
|
from collections import OrderedDict
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from sched.aggregator import average_weights
|
||
|
|
|
||
|
|
|
||
|
|
def _make_state(val: float) -> OrderedDict:
|
||
|
|
return OrderedDict(
|
||
|
|
weight=torch.full((3, 4), val),
|
||
|
|
bias=torch.full((3,), val),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_average_identical():
|
||
|
|
w = _make_state(1.0)
|
||
|
|
result = average_weights([w, w])
|
||
|
|
for key in w:
|
||
|
|
assert torch.allclose(result[key], w[key])
|
||
|
|
|
||
|
|
|
||
|
|
def test_average_different():
|
||
|
|
a = _make_state(0.0)
|
||
|
|
b = _make_state(2.0)
|
||
|
|
result = average_weights([a, b])
|
||
|
|
expected = _make_state(1.0)
|
||
|
|
for key in expected:
|
||
|
|
assert torch.allclose(result[key], expected[key])
|
||
|
|
|
||
|
|
|
||
|
|
def test_average_three():
|
||
|
|
a = _make_state(0.0)
|
||
|
|
b = _make_state(3.0)
|
||
|
|
c = _make_state(6.0)
|
||
|
|
result = average_weights([a, b, c])
|
||
|
|
expected = _make_state(3.0)
|
||
|
|
for key in expected:
|
||
|
|
assert torch.allclose(result[key], expected[key])
|
||
|
|
|
||
|
|
|
||
|
|
def test_single():
|
||
|
|
w = _make_state(5.0)
|
||
|
|
result = average_weights([w])
|
||
|
|
for key in w:
|
||
|
|
assert torch.allclose(result[key], w[key])
|
||
|
|
|
||
|
|
|
||
|
|
def test_empty_raises():
|
||
|
|
with pytest.raises(ValueError):
|
||
|
|
average_weights([])
|
||
|
|
|
||
|
|
|
||
|
|
def test_preserves_keys():
|
||
|
|
w = OrderedDict(fc1_weight=torch.ones(2, 2), fc1_bias=torch.zeros(2))
|
||
|
|
result = average_weights([w, w])
|
||
|
|
assert list(result.keys()) == ["fc1_weight", "fc1_bias"]
|