def _test_distrib_integration(device):
import numpy as np
from ignite.engine import Engine
rank = idist.get_rank()
n_iters = 80
s = 50
offset = n_iters * s
y_true = torch.arange(0, offset * idist.get_world_size(), dtype=torch.float).to(device)
y_preds = torch.ones(offset * idist.get_world_size(), dtype=torch.float).to(device)
def update(engine, i):
return (
y_preds[i * s + offset * rank : (i + 1) * s + offset * rank],
y_true[i * s + offset * rank : (i + 1) * s + offset * rank],
)
engine = Engine(update)
m = MeanAbsoluteError()
m.attach(engine, "mae")
data = list(range(n_iters))
engine.run(data=data, max_epochs=1)
assert "mae" in engine.state.metrics
res = engine.state.metrics["mae"]
true_res = np.mean(np.abs((y_true - y_preds).cpu().numpy()))
assert pytest.approx(res) == true_res
def _test(metric_device):
engine = Engine(update)
m = MeanAbsoluteError(device=metric_device)
m.attach(engine, "mae")
data = list(range(n_iters))
engine.run(data=data, max_epochs=1)
assert "mae" in engine.state.metrics
res = engine.state.metrics["mae"]
true_res = np.mean(np.abs((y_true - y_preds).cpu().numpy()))
assert pytest.approx(res) == true_res
