Repository: pytorch
Issue: LeakyReLU and Elu use more VRAM than needed
Labels: module: autograd,module: nn,triaged

## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
Both LeakyReLU and Elu need about twice the VRAM for storing tensors for the backward-passes.

This results from the out-of-place versions storing the input-tensor for the gradient calculation with autograd, where saving the output-tensor would be valid as well. However, the output-tensors are often already saved by the following layer, so that only one instead of two tensors would have to be saved when saving only the output-tensor, similar to ReLU #63027.

However, this needs an adaptation of the autograd generation to allow some saved-tensors to be set at run-time.

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
LeakyReLU is a standard activation function and used by many nets, where especially for convolutional networks the size of data-tensors is mostly restricting the network size regarding the VRAM, so that needing half the VRAM for activation functions enables the use of larger networks.

The in-place version of LeakyReLU and Elu does already use the output-tensor instead of the input-tensor. However, they need to restrict other input parameter e.g. for LeakyReLU `negative_slope >= 0`. With out-of-place versions being the standard used activation functions in PyTorch, many users are using the more VRAM needy version, where it also inconsistent for out-of-place and in-place versions using different restrictions needing a differing amount of VRAM.

## Pitch

<!-- A clear and concise description of what you want to happen. -->
For the new version, both LeakyReLU and Elu save for e.g. `negative_slope >= 0` for LeakyReLU the output-tensor `result` for gradient calculations instead of the input-tensor `self`. For e.g. `negative_slope < 0`, the input-tensor `self` has to be saved, due being ambiguous regarding the gradient calculations #63027.

There are multiple ways to realize that with `derivatives.yaml`. However, so far saved tensors are determined at compile-time, so that `gen_autograd.py` would need to be improved to allow setting some tensors to be saved at run-time depending on the function parameters e.g. `save_tensor = negative_slope.to<double>() >= 0 ? result : self;`.

### Option 1
Determine saved tensor sets for both `negative_slope >= 0` and `negative_slope < 0` at compile time, and save only the tensors of the set corresponding to whether `negative_slope` is larger or smaller than 0.

New `derivatives.yaml`:
```yaml
- name: leaky_relu(Tensor self, Scalar negative_slope=0.01) -> Tensor
  self: "negative_slope.to<double>() >= 0 ? leaky_relu_backward(grad, result, negative_slope, true) else leaky_relu_backward(grad, self, negative_slope, false)"
  result: auto_element_wise
```

### Option 2:
Use a new variable similar to how `return` works, where the variable is set during run-time.

New `derivatives.yaml`:
```yaml
- name: leaky_relu(Tensor self, Scalar negative_slope=0.01) -> Tensor
  save: "negative_slope.to<double>() >= 0 ? result : self"
  self: leaky_relu_backward(grad, save, negative_slope, false)
  result: auto_element_wise
```

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
Do nothing. However, this uses twice the VRAM for both LeakyReLU and Elu.

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->
The following script shows an implementation of the requested functionality for LeakyReLU when using torch.autograd.Function. It shows the time needed for forward + backward propagation in a closed-loop set-up of 10 iterations, and the VRAM needed without the input-tensors and network parameter.

Its results show that about half the VRAM is needed for `negative_slope >= 0`, and the same VRAM for `negative_slope < 0`.

``` python
import torch as th
from torch import nn
import time

from torch.nn import functional
from torch.autograd import Function

ITERATION = 10
DEVICE = 0
NEGATIVE_SLOPE = 0.01
SHAPE = [1, 256, 512, 512]


class LeakyReluFunction(Function):
    @staticmethod
    def forward(ctx, input: th.Tensor, negative_slope=1e-2) -> th.Tensor:
        relu = functional.leaky_relu(input, negative_slope)
        ctx.save_for_backward(relu if negative_slope >= 0 else input)
        ctx.negative_slope = negative_slope
        return relu

    @staticmethod
    def backward(ctx, grad: th.Tensor):
        relu, = ctx.saved_tensors
        negative_slope = ctx.negative_slope
        grad *= negative_slope * (relu <= 0) + (relu > 0)
        return grad, None


class LeakyReLU(nn.Module):
    def __init__(self, negative_slope: float = 1e-2):
        super(LeakyReLU, self).__init__()
        self.negative_slope = negative_slope

    def forward(self, input: th.Tensor) -> th.Tensor:
        return LeakyReluFunction.apply(input, self.negative_slope)


def speed_test(model: nn.Module, module: nn.Module, input: th.Tensor, output: th.Tensor, criterion, iteration=10):
    device = input.device
    module = module.to(device=device)

    input.grad = None
    input.requires_grad_()

    model = model.to(device=device)

    th.cuda.synchronize()
    th.cuda.reset_peak_memory_stats()
    vram_0 = th.cuda.max_memory_allocated()
    time_0 = time.time()

    prediction = input
    for _ in range(iteration):
        prediction = model(prediction)
        prediction = module(prediction)

    th.cuda.synchronize()
    vram_forward = th.cuda.max_memory_allocated()
    loss = criterion(prediction, output)
    loss.backward()

    th.cuda.synchronize()
    print(f'Forward + Backward: {time.time() - time_0:1.2f}s, \n'
          f'Forward (VRAM):     {(vram_forward - vram_0)/1e9:1.2f}GB, \n'
          f'Backward (VRAM):    {(th.cuda.max_memory_allocated() - vram_forward) / 1e9:1.2f}GB')

    return prediction, input.grad


if __name__ == "__main__":
    device = th.device(DEVICE)

    input = th.rand(size=SHAPE, device=device)
    label = th.rand_like(input)

    model = nn.Conv2d(
        in_channels=input.shape[1],
        out_channels=input.shape[1],
        kernel_size=(3, 3),
        padding=(1, 1),
        bias=True).to(device=device)

    criterion = nn.MSELoss()

    # For initializing device
    for _ in range(10):
        criterion(model(input), label).backward()

    print('New LeakyReLU:')
    pr_new = speed_test(model, LeakyReLU(negative_slope=NEGATIVE_SLOPE), input, label, criterion, ITERATION)

    print('')
    print('LeakyReLU:')
    pr_relu = speed_test(model, nn.LeakyReLU(negative_slope=NEGATIVE_SLOPE), input, label, criterion, ITERATION)

    print('')
    print(f'Max abs difference prediction: {th.max((pr_new[0] - pr_relu[0]).abs()).detach().cpu():1.2e}')
    print(f'Max abs difference grad:       {th.max((pr_new[1] - pr_relu[1]).abs()).detach().cpu():1.2e}')

```

### Output `negative_slope >= 0`
New LeakyReLU:
Forward + Backward: 0.56s, 
Forward (VRAM):     2.95GB, 
Backward (VRAM):    1.49GB

LeakyReLU:
Forward + Backward: 0.54s, 
Forward (VRAM):     5.37GB, 
Backward (VRAM):    1.49GB

Max abs difference prediction: 0.00e+00
Max abs difference grad:       0.00e+00

### Output `negative_slope < 0`
New LeakyReLU:
Forward + Backward: 0.59s, 
Forward (VRAM):     5.37GB, 
Backward (VRAM):    1.49GB

LeakyReLU:
Forward + Backward: 0.53s, 
Forward (VRAM):     5.37GB, 
Backward (VRAM):    1.49GB

Max abs difference prediction: 0.00e+00
Max abs difference grad:       0.00e+00

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @mruberry @jbschlosser