Pruning callback causes GPU memory leak when used iteratively

[MohammedAljahdali]

Hi, I have a script that does the following logic:

model = Model()
dm = DataModule()
callback_a = CallbackA()
for _ in range(N):
     callbacks = [CallbackB(), CallbackC(), callback_a]
     trainer = Trainer(callbacks=callbacks, ...)
     trainer.fit(model, dm)
     trainer.test()

There is much more going on, but to keep it simple this is the flow that I have, my question is there another preferable way to what I just did? Also, in my code I have some memory leak, that happens after each loop iteration, could this be somehow related to the trainer object, not being deleted properly?

[MohammedAljahdali]

I found out about this issue https://discuss.pytorch.org/t/how-to-avoid-memory-leak-when-training-multiple-models-sequentially/100315 which I think confirms that the memory leak is related to Lightning, but the question is how can we deal with it?

[MohammedAljahdali]

This the repo that contain my code: https://github.com/MohammedAljahdali/shrinkbench/tree/NAS, to use the following to reproduce the memory leak python run.py experiment=lth_02_20_cifar100_finetune_darts_sgd_88.yaml model.module.finetune_darts.n_layers=16 logger.wandb.project=test_runs trainer.max_epochs=2 +trainer.limit_train_batches=2 +trainer.limit_val_batches=1 +trainer.limit_test_batches=1 dataset=cifar10
This code require that you use wandb, and you can specify the name of your project in the args.

[carmocca]

Can you provide more details about the alleged memory leaks?

You could try and see if putting gc.collect() inside the start of your loop helps to delete any leftover memory from the previous iteration

[MohammedAljahdali]

Tried torch.cuda.empty_cache() it did not work.

[carmocca]

Can you try to reproduce it using our bug report model?

https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pl_examples/bug_report_model.py

[MohammedAljahdali]

So I managed to reproduce it, and it turns out that the memory leak is related to the PyTorch/pytorch_lightning pruning. It seems that each time pruning is applied the GPU memory usage increase. In my case, I prune at the end of each training, and that is why it seemed like the trainer causing this issue. However, I also tried to use the default lightning pruning callback, and the memory seems to increase after each epoch.

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 = MyModelPruning(
    #         n_levels=N,
    #         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
    #     )
    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()

If you want to use my pruning callback comment it off in the above code, and add the following code:

_PARAM_TUPLE = Tuple[nn.Module, str]
_PARAM_LIST = Union[List[_PARAM_TUPLE], Tuple[_PARAM_TUPLE]]


def flops(model, input):
    """Compute Multiply-add FLOPs estimate from model

    Arguments:
        model {torch.nn.Module} -- Module to compute flops for
        input {torch.Tensor} -- Input tensor needed for activations

    Returns:
        tuple:
        - int - Number of total FLOPs
        - int - Number of FLOPs related to nonzero parameters
    """
    FLOP_fn = {
        nn.Conv2d: _conv2d_flops,
        nn.Linear: _linear_flops,
    }

    total_flops = nonzero_flops = 0
    activations = get_activations(model, input)

    # The ones we need for backprop
    for m, (act, _) in activations.items():
        if m.__class__ in FLOP_fn:
            w = m.weight.detach().cpu().numpy().copy()
            module_flops = FLOP_fn[m.__class__](m, act)
            total_flops += module_flops
            # For our operations, all weights are symmetric so we can just
            # do simple rule of three for the estimation
            nonzero_flops += module_flops * nonzero(w).sum() / np.prod(w.shape)

    return total_flops, nonzero_flops


def nonzero(tensor):
    """Returns absolute number of values different from 0

    Arguments:
        tensor {numpy.ndarray} -- Array to compute over

    Returns:
        int -- Number of nonzero elements
    """
    return np.sum(tensor != 0.0)


def dense_flops(in_neurons, out_neurons):
    """Compute the number of multiply-adds used by a Dense (Linear) layer"""
    return in_neurons * out_neurons


def _conv2d_flops(module, activation):
    # Auxiliary func to use abstract flop computation

    # Drop batch & channels. Channels can be dropped since
    # unlike shape they have to match to in_channels
    input_shape = activation.shape[2:]
    # TODO Add support for dilation and padding size
    return conv2d_flops(in_channels=module.in_channels,
                        out_channels=module.out_channels,
                        input_shape=input_shape,
                        kernel_shape=module.kernel_size,
                        padding=module.padding,
                        strides=module.stride,
                        dilation=module.dilation)


def _linear_flops(module, activation):
    # Auxiliary func to use abstract flop computation
    return dense_flops(module.in_features, module.out_features)


def conv2d_flops(in_channels, out_channels, input_shape, kernel_shape,
                 padding='same', strides=1, dilation=1):
    """Compute the number of multiply-adds used by a Conv2D layer
    """
    # validate + sanitize input
    assert in_channels > 0
    assert out_channels > 0
    assert len(input_shape) == 2
    assert len(kernel_shape) == 2
    # padding = padding.lower()
    padding = padding if type(padding) is tuple else (padding, padding)
    try:
        strides = tuple(strides)
    except TypeError:
        # if one number provided, make it a 2-tuple
        strides = (strides, strides)
    dilation = dilation if type(dilation) is tuple else (dilation, dilation)

    output_dims = [0, 0]
    output_dims[0] = (input_shape[0] + 2 * padding[0] -
                      (dilation[0] * (kernel_shape[0] - 1) + 1)) // strides[0] + 1
    output_dims[1] = (input_shape[1] + 2 * padding[1] -
                      (dilation[1] * (kernel_shape[1] - 1) + 1)) // strides[1] + 1

    # output_shape = (int(out_nrows), int(out_ncols))

    # work to compute one output spatial position
    nflops = in_channels * out_channels * int(np.prod(kernel_shape))

    # total work = work per output position * number of output positions
    return nflops * int(np.prod(output_dims))


def get_activations(model, input):
    activations = OrderedDict()

    def store_activations(module, input, output):
        if isinstance(module, nn.ReLU):
            # TODO ResNet18 implementation reuses a
            # single ReLU layer?
            return
        assert module not in activations, \
            f"{module} already in activations"
        # TODO [0] means first input, not all models have a single input
        activations[module] = (input[0].detach().cpu().numpy().copy(),
                               output.detach().cpu().numpy().copy(),)

    fn, hooks = hook_applyfn(store_activations, model, forward=True)
    model.apply(fn)
    with torch.no_grad():
        model(input)

    for h in hooks:
        h.remove()

    return activations


def hook_applyfn(hook, model, forward=False, backward=False):
    """

    [description]

    Arguments:
        hook {[type]} -- [description]
        model {[type]} -- [description]

    Keyword Arguments:
        forward {bool} -- [description] (default: {False})
        backward {bool} -- [description] (default: {False})

    Returns:
        [type] -- [description]
    """
    assert forward ^ backward, \
        "Either forward or backward must be True"
    hooks = []

    def register_hook(module):
        if (
                not isinstance(module, nn.Sequential)
                and
                not isinstance(module, nn.ModuleList)
                and
                not isinstance(module, nn.ModuleDict)
                and
                not (module == model)
        ):
            if forward:
                hooks.append(module.register_forward_hook(hook))
            if backward:
                hooks.append(module.register_backward_hook(hook))

    return register_hook, hooks

class MyModelPruning(pl.callbacks.ModelPruning):

    def __init__(self, n_levels: int, **kwargs):
        super().__init__(**kwargs)
        self.level = ''
        self.n_levels = n_levels
        self.logger = None
        self.module: Optional[nn.Module] = None
        self.x: Optional[torch.T] = None

    def filter_parameters_to_prune(self, parameters_to_prune: Optional[_PARAM_LIST] = None):
        return [module for module in parameters_to_prune if self.can_prune(module[0])]

    def can_prune(self, module):
        if hasattr(module, 'is_classifier'):
            return not module.is_classifier
        return True

    def on_train_epoch_end(self, trainer, pl_module: LightningModule):
        pass

    # def on_test_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
    #     pass

    def on_test_end(self, trainer, pl_module):

        current_epoch = trainer.current_epoch

        prune = self._apply_pruning(current_epoch) if isinstance(self._apply_pruning, Callable) else self._apply_pruning
        amount = self.amount(current_epoch) if isinstance(self.amount, Callable) else self.amount
        if not prune or not amount:
            return
        self.level = pl_module.run_id
        if str(self.level) == str(self.n_levels - 1):
            return
        # This avoid pruning the model again when performing the second test to see the model results after pruning
        if pl_module.test_counter % 2 != 0:
            pl_module.test_counter += 1
            return
        pl_module.test_counter += 1
        # amount = 1
        # for _ in range(int(self.level)+1):
        #     amount *= 0.5

        # TODO: Make it general to support multiple loggers
        self.logger = pl_module.logger
        self.module = pl_module.net
        self.x, y = next(iter(trainer.datamodule.train_dataloader()))
        self.x = self.x.to(pl_module.device)
        self.apply_pruning(amount)

        if (
                self._use_lottery_ticket_hypothesis(current_epoch)
                if isinstance(self._use_lottery_ticket_hypothesis, Callable) else self._use_lottery_ticket_hypothesis
        ):
            self.apply_lottery_ticket_hypothesis()

    @rank_zero_only
    def _log_sparsity_stats(
            self, prev: List[Tuple[int, int]], curr: List[Tuple[int, int]], amount: Union[int, float] = 0
    ):
        total_params = sum(p.numel() for layer, _ in self._parameters_to_prune for p in layer.parameters())
        prev_total_zeros = sum(zeros for zeros, _ in prev)
        curr_total_zeros = sum(zeros for zeros, _ in curr)
        self.logger.experiment[0].summary[f"level-{self.level}/prev_size"] = total_params - prev_total_zeros
        self.logger.experiment[0].summary[f"level-{self.level}/size"] = total_params - curr_total_zeros
        self.logger.experiment[0].summary[f"level-{self.level}/pruned"] = curr_total_zeros / total_params

        self.logger.experiment[0].summary[f"level-{self.level}/compression_ratio"] = total_params / (
                total_params - curr_total_zeros)

        # FLOPS

        ops, ops_nz = flops(self.module, self.x)
        self.logger.experiment[0].summary[f"level-{self.level}/flops"] = ops
        self.logger.experiment[0].summary[f"level-{self.level}/flops_nz"] = ops_nz
        self.logger.experiment[0].summary[f"level-{self.level}/theoretical_speedup"] = ops / ops_nz

        if self._verbose == 2:
            for i, (module, name) in enumerate(self._parameters_to_prune):
                prev_mask_zeros, prev_mask_size = prev[i]
                curr_mask_zeros, curr_mask_size = curr[i]

[MohammedAljahdali]

Should we move this dissuction into an issue?

[carmocca]

Yes