diff --git a/tests/test_pid_grad_buffer.py b/tests/test_pid_grad_buffer.py new file mode 100644 index 00000000..ecdbdd2f --- /dev/null +++ b/tests/test_pid_grad_buffer.py @@ -0,0 +1,28 @@ +"""The first PID step must remember its gradient.""" +import torch +import torch_optimizer as optim + + +def test_second_step_uses_the_previous_gradient(): + param = torch.nn.Parameter(torch.tensor([0.0])) + opt = optim.PID( + [param], + lr=1.0, + momentum=0.5, + dampening=0.0, + weight_decay=0.0, + integral=0.0, + derivative=10.0, + ) + param.grad = torch.tensor([1.0]) + opt.step() + assert torch.equal( + opt.state[param]["grad_buffer"], torch.tensor([1.0]) + ) + # first derivative is 0: there is no earlier gradient + assert torch.equal(param.detach(), torch.tensor([-1.0])) + + param.grad = torch.tensor([3.0]) + opt.step() + # D = (1 - momentum) * (3 - 1) = 1, so the step is 3 + 10 * 1 + assert torch.allclose(param.detach(), torch.tensor([-14.0])) diff --git a/torch_optimizer/pid.py b/torch_optimizer/pid.py index 44dad253..ca6bf1c4 100644 --- a/torch_optimizer/pid.py +++ b/torch_optimizer/pid.py @@ -104,7 +104,10 @@ def step(self, closure: OptLossClosure = None) -> OptFloat: g_buf = param_state["grad_buffer"] = torch.zeros_like( p, memory_format=torch.preserve_format ) - g_buf = d_p + # copy into the buffer. Rebinding g_buf = d_p left + # the stored gradient at 0, so the next step's + # derivative was taken against a zero previous grad. + g_buf.copy_(d_p) d_buf = param_state["d_buffer"] = torch.zeros_like( p, memory_format=torch.preserve_format