| | |
| | |
| | |
| | |
| | |
| |
|
| | import torch |
| |
|
| |
|
| | def check_require_grad(t): |
| | return isinstance(t, torch.Tensor) and t.requires_grad |
| |
|
| |
|
| | class CheckpointFunction(torch.autograd.Function): |
| | @staticmethod |
| | def forward(ctx, run_function, length, *args): |
| | ctx.run_function = run_function |
| | ctx.input_tensors = list(args[:length]) |
| | ctx.input_params = list(args[length:]) |
| | with torch.no_grad(): |
| | output_tensors = ctx.run_function(*ctx.input_tensors) |
| | return output_tensors |
| |
|
| | @staticmethod |
| | def backward(ctx, *output_grads): |
| | for i in range(len(ctx.input_tensors)): |
| | temp = ctx.input_tensors[i] |
| | if check_require_grad(temp): |
| | ctx.input_tensors[i] = temp.detach() |
| | ctx.input_tensors[i].requires_grad = temp.requires_grad |
| | with torch.enable_grad(): |
| | output_tensors = ctx.run_function(*ctx.input_tensors) |
| | to_autograd = list(filter(check_require_grad, ctx.input_tensors)) |
| | output_tensors, output_grads = zip(*filter(lambda t: t[0].requires_grad, zip(output_tensors, output_grads))) |
| | input_grads = torch.autograd.grad(output_tensors, to_autograd + ctx.input_params, output_grads, allow_unused=True) |
| | input_grads = list(input_grads) |
| | for i in range(len(ctx.input_tensors)): |
| | if not check_require_grad(ctx.input_tensors[i]): |
| | input_grads.insert(i, None) |
| | return (None, None) + tuple(input_grads) |
| |
|
