tinynn/prune/admm_pruner.py

import os

import torch
import torch.distributed as dist

from tinynn.util.util import get_logger
from tinynn.prune.oneshot_pruner import OneShotChannelPruner

log = get_logger(__name__)


class ADMMPruner(OneShotChannelPruner):
    required_params = ('sparsity', 'metrics', 'admm_iterations', 'admm_epoch', 'rho', 'admm_lr')
    required_context_params = ('val_loader', 'train_loader', 'train_func', 'validate_func', 'optimizer', 'criterion')
    default_values = {'admm_save_freq': 1, 'admm_valid_freq': 1, 'admm_dir': 'admm_train/'}
    context_from_params_dict = {
        'optimizer': ['admm_optimizer', 'optimizer'],
        'criterion': ['admm_criterion', 'criterion'],
    }
    condition_dict = {
        'admm_iterations': lambda x: 0 < x,
        'admm_epoch': lambda x: 0 < x,
        'rho': lambda x: 0 < x < 1,
        'admm_lr': lambda x: 0 < x < 1,
    }

    admm_iterations: int
    admm_epoch: int
    rho: float
    admm_lr: float
    admm_dir: str
    admm_save_freq: int
    admm_valid_freq: int

    def __init__(self, model, dummy_input, config, context):
        super().__init__(model, dummy_input, config)

        self.parse_context(context)

        self.Z = {}
        self.U = {}

    def prune(self):
        """The main function for pruning"""

        log.info('Start ADMM training')
        old_criterion = self.context.criterion
        self.register_mask()
        for iteration in range(1, self.admm_iterations + 1):
            for epoch in range(1, self.admm_epoch + 1):
                self.context.epoch = epoch + (iteration - 1) * self.admm_epoch
                self.adjust_learning_rate()
                self.context.criterion = self.construct_admm_criterion(old_criterion)
                if dist.is_available() and dist.is_initialized():
                    self.context.train_loader.sampler.set_epoch(self.context.epoch)
                self.context.train_func(self.model, self.context)

                if not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0:
                    if self.context.epoch % self.admm_save_freq == 0:
                        save_path = os.path.join(self.admm_dir, f'epoch_{self.context.epoch}.pth')
                        os.makedirs(os.path.dirname(save_path), exist_ok=True)
                        log.info("Saving model to {}".format(save_path))
                        torch.save(self.model.state_dict(), save_path)

                    if self.context.validate_func is not None and self.context.epoch % self.admm_valid_freq == 0:
                        # According to https://github.com/pytorch/pytorch/issues/54059, when validating via DDP,
                        # it needs to be done on the original module.
                        if dist.is_available() and dist.is_initialized():
                            self.context.validate_func(self.model.module, self.context)
                        else:
                            self.context.validate_func(self.model, self.context)

                self.context.best_epoch = self.context.epoch

            self.admm_params_update()

        self.context.criterion = old_criterion
        # Need to reset the masker before final pruning
        self.graph_modifier.reset_masker()
        super().prune()

    def adjust_learning_rate(self):
        epoch = self.context.epoch
        admm_epoch = self.admm_epoch
        if (epoch - 1) % admm_epoch == 0:
            lr = self.admm_lr
        else:
            admm_epoch_offset = (epoch - 1) % admm_epoch
            # LR is updated roughly every 1/3 admm_epoch.
            admm_step = admm_epoch / 3
            lr = self.admm_lr * (0.1 ** (admm_epoch_offset // admm_step))

        for param_group in self.context.optimizer.param_groups:
            param_group['lr'] = lr

    def construct_admm_criterion(self, old_criterion):
        def criterion_func(output, target):
            loss = old_criterion(output, target)
            for n in self.center_nodes:
                if n.unique_name not in self.Z or n.unique_name not in self.U:
                    continue

                loss += (
                    0.5
                    * self.rho
                    * (torch.norm(n.module.weight - self.Z[n.unique_name] + self.U[n.unique_name], p=2) ** 2)
                )
            return loss

        return criterion_func

    def register_mask(self):
        super().register_mask()

        for m in self.graph_modifier.modifiers.values():
            # Disable mask here so that we will use them to update U and Z.
            # They won't be applied so that the training process won't be affected.
            m.disable_mask()
            if m.node in self.center_nodes and m.dim_changes_info.pruned_idx_o:
                device = m.module().weight.data.device
                self.Z[m.unique_name()] = m.module().weight.detach() * m.masker().get_mask('weight').to(device=device)
                self.U[m.unique_name()] = torch.zeros_like(self.Z[m.unique_name()], device=device)

    def admm_params_update(self):
        self.graph_modifier.reset_masker()
        importance = {}
        for sub_graph in self.graph_modifier.sub_graphs.values():
            for m in sub_graph.modifiers:
                if m.node in self.center_nodes and m in sub_graph.dependent_centers:
                    if m.unique_name() not in self.U:
                        continue
                    self.Z[m.unique_name()] = (m.module().weight + self.U[m.unique_name()]).detach()
                    importance[m.unique_name()] = self.metric_func(self.Z[m.unique_name()], m.module())

            sub_graph.calc_prune_idx(importance, self.sparsity)
            log.info(f"subgraph [{sub_graph.center}] compute over")

        for m in self.graph_modifier.modifiers.values():
            m.register_mask(self.graph_modifier.modifiers, importance, self.sparsity)

            # Disable mask here so that we will use them to update U and Z.
            # They won't be applied so that the training process won't be affected.
            m.disable_mask()
            if m.node in self.center_nodes and m.dim_changes_info.pruned_idx_o:
                weight = m.module().weight
                self.Z[m.unique_name()] = self.Z[m.unique_name()] * m.masker().get_mask('weight')
                self.U[m.unique_name()] = weight - self.Z[m.unique_name()] + self.U[m.unique_name()]

        # Sync ADMM parameters
        if dist.is_available() and dist.is_initialized():
            for state in (self.U, self.Z):
                for param in state.values():
                    dist.broadcast(param, 0)