Pruning callback causes GPU memory leak when used iteratively

@tchaton is already working on this.

Since Aug 27, 2021.

### Discussed in https://github.com/PyTorchLightning/pytorch-lightning/discussions/8363

After trying to reproduce the issue with the boring model, it turned out that the cause of the momery leak is not the reinitialization of the trainer, but the pruning callback itself.

Dependices:
```
- python=3.8
- pip
- cudatoolkit=10.2
- pytorch=1.8.1
- torchvision=0.9.1
- pytorch-lightning>=1.3.2
- torchmetrics>=0.3.2
- wandb>=0.10.30
```

Code to reproduce:
```
import os
import gc

import torch
from torch.utils.data import DataLoader, Dataset
from torch import nn
import wandb
from pytorch_lightning import LightningModule, Trainer, LightningDataModule
from pytorch_lightning.loggers import WandbLogger
import torchvision
import pytorch_lightning as pl
import numpy as np
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
from pytorch_lightning.utilities.distributed import rank_zero_only
from collections import OrderedDict, defaultdict

class DM(LightningDataModule):

def __init__(self, n_features, n_samples, batch_size):
super().__init__()
self.train_dataset = torchvision.datasets.CIFAR10(root='/tmp', download=True,
transform=torchvision.transforms.ToTensor())
self.val_dataset = torchvision.datasets.CIFAR10(root='/tmp', transform=torchvision.transforms.ToTensor())
self.test_dataset = torchvision.datasets.CIFAR10(root='/tmp', transform=torchvision.transforms.ToTensor())
self.batch_size = batch_size

def setup(self, stage=None):
pass

def train_dataloader(self):
return DataLoader(
dataset=self.train_dataset,
batch_size=self.batch_size,
shuffle=True
)

def val_dataloader(self):
return DataLoader(
dataset=self.val_dataset,
batch_size=self.batch_size,
shuffle=False
)

def test_dataloader(self):
return DataLoader(
dataset=self.test_dataset,
batch_size=self.batch_size,
shuffle=False
)

class BoringModel(LightningModule):

def __init__(self, in_features):
super().__init__()
self.run_id = 0
self.test_counter = 0
self.net = torchvision.models.resnet18()
self.net.fc.is_classifier = True
self.loss = torch.nn.CrossEntropyLoss()

# self.layer = torch.nn.Sequential(
# torch.nn.Linear(in_features, in_features // 2),
# torch.nn.ReLU(),
# torch.nn.Linear(in_features // 2, in_features // 4),
# torch.nn.ReLU(),
# torch.nn.Linear(in_features // 4, in_features // 8),
# torch.nn.ReLU(),
# torch.nn.Linear(in_features // 8, 2),
# )

def forward(self, x):
return self.net(x)

def training_step(self, batch, batch_idx):
loss = self.loss(self(batch[0]), batch[1])
self.log(f"train_loss {self.run_id}", loss)
return {"loss": loss}

def validation_step(self, batch, batch_idx):
loss = self.loss(self(batch[0]), batch[1])
self.log(f"valid_loss {self.run_id}", loss)

def test_step(self, batch, batch_idx):
loss = self.loss(self(batch[0]), batch[1])
self.log(f"test_loss {self.run_id}", loss)

def configure_optimizers(self):
return torch.optim.SGD(self.net.parameters(), lr=1)

def run():
in_features = 224
N = 20
dm = DM(n_samples=40, n_features=in_features, batch_size=256)
model = BoringModel(in_features=in_features)
loggers = [WandbLogger(project='test_runs', save_dir='/tmp')]
pruning_callback = pl.callbacks.ModelPruning(
apply_pruning=True, use_lottery_ticket_hypothesis=True,
pruning_fn='l1_unstructured', use_global_unstructured=True, verbose=1, make_pruning_permanent=False,
amount=0.5
)

for i in range(N):
callbacks = [
pl.callbacks.ModelCheckpoint(monitor=f"valid_loss {model.run_id}", save_top_k=1, save_last=True),
pl.callbacks.EarlyStopping(monitor=f"valid_loss {model.run_id}", )
]
callbacks.append(pruning_callback)
trainer = Trainer(
default_root_dir='/tmp',
limit_train_batches=4,
limit_val_batches=4,
limit_test_batches=4,
num_sanity_val_steps=0,
max_epochs=5,
logger=loggers,
gpus=1,
callbacks=callbacks
)

trainer.fit(model, datamodule=dm)
trainer.test(model, ckpt_path='best')
trainer.test(model)
model.run_id += 1
counter = 0
for obj in gc.get_objects():
try:
if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
counter += 1
except:
pass
print(f"Number of tensors after level {i - 1} is {counter}")

if __name__ == '__main__':
run()
```