def test_input_singleton(setup_rpc): class One(nn.Module): def __init__(self): super().__init__() self.fc = nn.Linear(1, 1) def forward(self, only_a): (a,) = only_a return (self.fc(a),) model = nn.Sequential(One()) model = Pipe(model, chunks=2) a = torch.rand(10, 1, requires_grad=True) (a_out,) = model((a,)).local_value() loss = a_out.mean() loss.backward() assert all(p.grad is not None for p in model.parameters()) assert a.grad is not None # Test with list a.grad = None for p in model.parameters(): p.grad = None (a_out,) = model([a]).local_value() loss = a_out.mean() loss.backward() assert all(p.grad is not None for p in model.parameters()) assert a.grad is not None
def test_simple_linears(): def sum_grad(parameters): return sum([p.grad.sum() for p in parameters if p.grad is not None]) def zero_grad(parameters): for p in parameters: p.grad = None inputs = torch.rand(8, 1) model = nn.Sequential(nn.Linear(1, 2), nn.Linear(2, 4), nn.Linear(4, 2), nn.Linear(2, 1),) # Without Pipe outputs = model(inputs) loss = outputs.mean() loss.backward() grad_without_pipe = sum_grad(model.parameters()) zero_grad(model.parameters()) # With Pipe model = Pipe(model, [2, 2], devices=["cpu", "cpu"], chunks=4) outputs = model(inputs) loss = outputs.mean() loss.backward() grad_with_pipe = sum_grad(model.parameters()) # Both grads should be identical. assert torch.allclose(grad_with_pipe, grad_without_pipe)
def test_input_singleton(): class One(nn.Module): def __init__(self): super().__init__() self.fc = nn.Linear(1, 1) def forward(self, only_a): (a,) = only_a return (self.fc(a),) model = nn.Sequential(One()) model = Pipe(model, balance=[1], devices=["cpu"], chunks=2) a = torch.rand(10, 1, requires_grad=True) (a_out,) = model((a,)) loss = a_out.mean() loss.backward() assert all(p.grad is not None for p in model.parameters()) assert a.grad is not None
def test_parameters(): model = nn.Sequential(nn.Linear(1, 1)) pipe = Pipe(model, balance=[1], devices=["cpu"], chunks=1) assert list(pipe.parameters()) != []
def test_parameters(): model = nn.Sequential(nn.Linear(1, 1)) pipe = Pipe(model, chunks=1) assert list(pipe.parameters()) != []