def _test_distrib_variable_accumulation(device):
mean_var = VariableAccumulation(lambda a, x: a + x, device=device)
y_true = torch.rand(100, device=device, dtype=torch.float64)
for y in y_true:
mean_var.update(y)
y_true = idist.all_reduce(y_true)
a, n = mean_var.compute()
assert a.item() == pytest.approx(y_true.sum().item())
assert n == len(y_true) * idist.get_world_size()
# check if call compute twice
a, n = mean_var.compute()
assert a.item() == pytest.approx(y_true.sum().item())
assert n == len(y_true) * idist.get_world_size()
mean_var = VariableAccumulation(lambda a, x: a + x, device=device)
y_true = torch.rand(50, 10, device=device, dtype=torch.float64)
for y in y_true:
mean_var.update(y)
y_true = idist.all_reduce(y_true)
a, n = mean_var.compute()
assert n == len(y_true) * idist.get_world_size()
np.testing.assert_almost_equal(a.cpu().numpy(),
y_true.sum(dim=0).cpu().numpy(),
decimal=4)
a, n = mean_var.compute()
assert n == len(y_true) * idist.get_world_size()
np.testing.assert_almost_equal(a.cpu().numpy(),
y_true.sum(dim=0).cpu().numpy(),
decimal=4)
def _test_apex_average(device, amp_mode, opt_level):
assert amp_mode == "apex"
assert device == "cuda"
model = Linear(1, 1)
if device:
model.to(device)
model.weight.data.zero_()
model.bias.data.zero_()
optimizer = SGD(model.parameters(), 0.1)
from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
mean_var = VariableAccumulation(lambda a, x: a + x)
y_true = torch.rand(100).float().to(device)
for y in y_true:
mean_var.update(y)
a, n = mean_var.compute()
assert a.item() == pytest.approx(y_true.sum().item())
assert n == len(y_true)
def test_variable_accumulation_wrong_inputs():
with pytest.raises(TypeError, match=r"Argument op should be a callable"):
VariableAccumulation(1)
with pytest.raises(TypeError, match=r"Output should be a number or torch.Tensor,"):
mean_acc = VariableAccumulation(lambda a, x: a + x)
mean_acc.update((1, 2))
with pytest.raises(TypeError, match=r"Output should be a number or torch.Tensor,"):
mean_acc = VariableAccumulation(lambda a, x: a + x)
mean_acc.update("a")
def _test_distrib_accumulator_device(device):
metric_devices = [torch.device("cpu")]
if device.type != "xla":
metric_devices.append(idist.device())
for metric_device in metric_devices:
m = VariableAccumulation(lambda a, x: x, device=metric_device)
assert m._device == metric_device
assert (
m.accumulator.device == metric_device
), f"{type(m.accumulator.device)}:{m.accumulator.device} vs {type(metric_device)}:{metric_device}"
m.update(torch.tensor(1, device=device))
assert (
m.accumulator.device == metric_device
), f"{type(m.accumulator.device)}:{m.accumulator.device} vs {type(metric_device)}:{metric_device}"
def test_variable_accumulation_mean_variable():
mean_var = VariableAccumulation(lambda a, x: a + x)
y_true = torch.rand(100)
for y in y_true:
mean_var.update(y)
a, n = mean_var.compute()
assert a.item() == pytest.approx(y_true.sum().item())
assert n == len(y_true)
mean_var = VariableAccumulation(lambda a, x: a + x)
y_true = torch.rand(100, 10)
for y in y_true:
mean_var.update(y)
a, n = mean_var.compute()
assert a.numpy() == pytest.approx(y_true.sum(dim=0).numpy())
assert n == len(y_true)
def test_variable_accumulation_mean_variable():
mean_var = VariableAccumulation(lambda a, x: a + x)
y_true = torch.rand(100)
for y in y_true:
mean_var.update(y)
a, n = mean_var.compute()
assert a.item() == pytest.approx(y_true.sum().item())
assert n == len(y_true)
mean_var = VariableAccumulation(lambda a, x: a + x)
y_true = torch.rand(100, 10)
for y in y_true:
mean_var.update(y)
a, n = mean_var.compute()
assert a.numpy() == pytest.approx(y_true.sum(dim=0).numpy())
assert n == len(y_true)
mean_var = VariableAccumulation(lambda a, x: a + x.sum(dim=0))
# iterate by batch of 16 samples
y_true = torch.rand(8, 16, 10)
for y in y_true:
mean_var.update(y)
a, n = mean_var.compute()
assert a.numpy() == pytest.approx(
y_true.reshape(-1, 10).sum(dim=0).numpy())
assert n == y_true.shape[0] * y_true.shape[1]