TensorRT

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 2021-11-09 13:38
# -------------------- sparse_masklib.py --------------------
import sys
import torch
import numpy as np
import collections
from itertools import permutations


def fill(x):
    return float(x.nonzero().size(0)) / torch.numel(x)


def reshape_1d(matrix, m):
    # If not a nice multiple of m, fill with zeroes.
    if matrix.shape[1] % m > 0:
        mat = torch.cuda.FloatTensor(
            matrix.shape[0], matrix.shape[1] + (m - matrix.shape[1] % m)
        ).fill_(0)
        mat[:, : matrix.shape[1]] = matrix
        shape = mat.shape
        return mat.view(-1, m), shape
    else:
        return matrix.view(-1, m), matrix.shape


valid_m4n2_1d_patterns = None


def compute_valid_1d_patterns(m, n):
    # Early exit if patterns was already created.
    global valid_m4n2_1d_patterns

    if m == 4 and n == 2 and valid_m4n2_1d_patterns is not None:
        return valid_m4n2_1d_patterns
    patterns = torch.zeros(m)
    patterns[:n] = 1
    valid_patterns = torch.Tensor(list(set(permutations(patterns.tolist()))))
    if m == 4 and n == 2:
        valid_m4n2_1d_patterns = valid_patterns
    return valid_patterns


""" m:n 1d structured best """


def mn_1d_best(matrix, m, n):
    # Find all possible patterns.
    patterns = compute_valid_1d_patterns(m, n).cuda()

    # Find the best m:n pattern (sum of non-masked weights).
    mask = torch.cuda.IntTensor(matrix.shape).fill_(1).view(-1, m)
    mat, shape = reshape_1d(matrix, m)
    pmax = torch.argmax(torch.matmul(mat.abs(), patterns.t()), dim=1)
    mask[:] = patterns[pmax[:]]
    mask = mask.view(matrix.shape)
    return mask


def m4n2_1d(mat, density):
    return mn_1d_best(mat, 4, 2)


def create_mask(tensor, pattern="m4n2_1d", density=0.5):
    # Reshape tensor and mask.
    shape = tensor.shape
    ttype = tensor.type()
    t = tensor.float().contiguous()

    # 1d-tensor
    if len(shape) == 1:
        t = t.view(1, shape[0])
        func = getattr(sys.modules[__name__], pattern, None)
        mask = func(t, density)
        return mask.view(shape).type(ttype)
    # 2d-tensor (in, out)
    elif len(shape) == 2:
        t = t.view(shape[0], shape[1])
        func = getattr(sys.modules[__name__], pattern, None)
        mask = func(t, density)
        return mask.view(shape).type(ttype)
    # 3d-tensor (batch, in, out)
    elif len(shape) == 3:
        t = t.view(shape[0] * shape[1], shape[2])
        func = getattr(sys.modules[__name__], pattern, None)
        mask = func(t, density)
        return mask.view(shape).type(ttype)
    # 4d-tensor (in, out, h, w)
    elif len(shape) == 4:
        # convs
        t = (
            t.permute(2, 3, 0, 1)
            .contiguous()
            .view(shape[2] * shape[3] * shape[0], shape[1])
        )
        func = getattr(sys.modules[__name__], pattern, None)
        mask = func(t, density)
        mask = (
            mask.view(shape[2], shape[3], shape[0], shape[1])
            .permute(2, 3, 0, 1)
            .contiguous()
        )
        return mask.view(shape).type(ttype)


# -------------------- asp.py --------------------
import types
import torch

torchvision_imported = True
try:
    import torchvision
except ImportError:
    print("[ASP][Warning] torchvision cannot be imported.")
    torchvision_imported = False


def eligible_modules(
    model, whitelist_layer_types, allowed_layer_names, disallowed_layer_names
):
    eligible_modules_list = []
    for name, mod in model.named_modules():
        if (
            isinstance(mod, whitelist_layer_types)
            and name not in disallowed_layer_names
        ):
            if allowed_layer_names is not None and name not in allowed_layer_names:
                continue
            eligible_modules_list.append((name, mod))
    return eligible_modules_list


class ASP:
    __model = None
    __verbosity = 0
    __optimizer = None
    __sparse_parameters = []
    __calculate_mask = None

    @classmethod
    def init_model_for_pruning(
        cls,
        model,
        mask_calculator="m4n2_1d",
        verbosity=3,
        whitelist=[torch.nn.Linear, torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d],
        allowed_layer_names=None,
        disallowed_layer_names=[],
        allow_recompute_mask=False,
        custom_layer_dict={},
    ):
        assert cls.__model is None, "ASP has been initialized already."
        cls.__model = model
        cls.__verbosity = verbosity

        if isinstance(mask_calculator, str):

            def create_mask_from_pattern(param):
                return create_mask(param, mask_calculator).bool()

            cls.__calculate_mask = create_mask_from_pattern
        else:
            cls.__calculate_mask = mask_calculator  # user defined function

        # function to extract variables that will be sparsified.
        # idea is that you will add one of these functions for each module type that can be sparsified.
        if torchvision_imported:
            print(
                "[ASP] torchvision is imported, can work with the MaskRCNN/KeypointRCNN from torchvision."
            )
            sparse_parameter_list = {
                torch.nn.Linear: ["weight"],
                torch.nn.Conv1d: ["weight"],
                torch.nn.Conv2d: ["weight"],
                torch.nn.Conv3d: ["weight"],
                torchvision.ops.misc.Conv2d: ["weight"],
            }
        else:
            sparse_parameter_list = {
                torch.nn.Linear: ["weight"],
                torch.nn.Conv1d: ["weight"],
                torch.nn.Conv2d: ["weight"],
                torch.nn.Conv3d: ["weight"],
            }
        if (
            custom_layer_dict
        ):  # Update default list to include user supplied custom (layer type : parameter tensor), make sure this tensor type is something ASP knows how to prune
            sparse_parameter_list.update(custom_layer_dict)
            whitelist += list(custom_layer_dict.keys())

        for module_type in whitelist:
            assert module_type in sparse_parameter_list, (
                "Module %s :: Don't know how to sparsify module." % module.dtype()
            )

        # find all sparse modules, extract sparse parameters and decorate
        def add_sparse_attributes(module_name, module):
            sparse_parameters = sparse_parameter_list[type(module)]
            for p_name, p in module.named_parameters():
                if p_name in sparse_parameters and p.requires_grad:
                    # check for NVIDIA's TC compatibility: we check along the horizontal direction
                    if p.dtype == torch.float32 and (
                        (p.size()[0] % 8) != 0 or (p.size()[1] % 16) != 0
                    ):  # User defines FP32 and APEX internally uses FP16 math
                        print(
                            "[ASP] Auto skipping pruning %s::%s of size=%s and type=%s for sparsity"
                            % (module_name, p_name, str(p.size()), str(p.dtype))
                        )
                        continue
                    if p.dtype == torch.float16 and (
                        (p.size()[0] % 8) != 0 or (p.size()[1] % 16) != 0
                    ):  # For Conv2d dim= K x CRS; we prune along C
                        print(
                            "[ASP] Auto skipping pruning %s::%s of size=%s and type=%s for sparsity"
                            % (module_name, p_name, str(p.size()), str(p.dtype))
                        )
                        continue

                    if cls.__verbosity >= 3:
                        print(
                            "[ASP] Sparsifying %s::%s of size=%s and type=%s for sparsity"
                            % (module_name, p_name, str(p.size()), str(p.dtype))
                        )

                    mask = torch.ones_like(p).bool()
                    buffname = p_name.split(".")[-1]  # buffer names cannot contain "."
                    module.register_buffer("__%s_mma_mask" % buffname, mask)
                    if allow_recompute_mask:
                        pruned = torch.zeros_like(p).cpu()
                        module.register_buffer("__%s_mma_pruned_p" % buffname, pruned)
                    else:
                        pruned = None
                    cls.__sparse_parameters.append(
                        (module_name, module, p_name, p, mask, pruned)
                    )
                else:
                    if cls.__verbosity >= 3:
                        print(
                            "[ASP] Not sparsifying %s::%s of size=%s and type=%s"
                            % (module_name, p_name, str(p.size()), str(p.dtype))
                        )

        for name, sparse_module in eligible_modules(
            model, tuple(whitelist), allowed_layer_names, disallowed_layer_names
        ):
            add_sparse_attributes(name, sparse_module)

    @classmethod
    def init_optimizer_for_pruning(cls, optimizer):
        assert cls.__optimizer is None, "ASP has initialized optimizer already."
        assert (
            cls.__calculate_mask is not None
        ), "Called ASP.init_optimizer_for_pruning before ASP.init_model_for_pruning."

        # store pointer to original optimizer step method
        cls.__optimizer = optimizer
        cls.__optimizer.__step = optimizer.step

        def __step(opt_self, *args, **kwargs):
            # prune gradients before step method
            with torch.no_grad():
                for (
                    module_name,
                    module,
                    p_name,
                    p,
                    mask,
                    pruned,
                ) in cls.__sparse_parameters:
                    if p.grad is not None:  # thx pjudd
                        p.grad.mul_(mask)
            # call original optimizer step method
            rval = opt_self.__step(*args, **kwargs)
            # prune parameters after step method
            with torch.no_grad():
                for (
                    module_name,
                    module,
                    p_name,
                    p,
                    mask,
                    pruned,
                ) in cls.__sparse_parameters:
                    p.mul_(mask)
            return rval

        cls.__optimizer.step = types.MethodType(__step, cls.__optimizer)

    @classmethod
    def compute_sparse_masks(cls):
        with torch.no_grad():
            for module_name, module, p_name, p, mask, pruned in cls.__sparse_parameters:
                if mask.sum() < mask.numel():  # when recalculating masks
                    # restore dense parameter if allow_recompute_mask is enabled
                    assert (
                        pruned is not None
                    ), "Unable to restore dense parameter because allow_recompute_mask == False"
                    p.add_(pruned.cuda())

                mask.set_(cls.__calculate_mask(p))

                if pruned is not None:  # stow away pruned weights to cpu
                    pruned.set_((p * (~mask)).cpu())

                p.mul_(
                    mask
                )  # in-place multiplication, so pruned weights are 0-values, hence checkpoint will have 0s for pruned weights
                if cls.__verbosity >= 2:
                    print(
                        "[ASP] Enabled %.2f%% sparsity for %s::%s of size=%s and type=%s"
                        % (
                            100.0 * mask.sum() / mask.numel(),
                            module_name,
                            p_name,
                            str(p.size()),
                            str(p.dtype),
                        )
                    )

    @classmethod
    def prune_trained_model(cls, model, optimizer):
        # add mask buffers to model (init_model_for_pruning), augment optimizer (init_optimizer_for_pruning) and compute masks (compute_sparse_masks)
        cls.init_model_for_pruning(
            model,
            mask_calculator="m4n2_1d",
            verbosity=2,
            whitelist=[torch.nn.Linear, torch.nn.Conv2d],
            allow_recompute_mask=False,
        )
        cls.init_optimizer_for_pruning(optimizer)
        cls.compute_sparse_masks()


# -------------------- test.py --------------------
import os
import numpy as np
import torch
from torch import nn
from torch import optim
from torch.utils.data import DataLoader, Dataset
import torch.nn.functional as F

model = None
optimizer = None


class ToyDataset(Dataset):
    def __init__(self):
        x = torch.round(torch.rand(1000) * 200)
        x = x.unsqueeze(1)
        x = torch.cat((x, x * 2, x * 3, x * 4, x * 5, x * 6, x * 7, x * 8), 1)
        self.X = x
        self.Y = self.X

    def __getitem__(self, index):
        return self.X[index], self.Y[index]

    def __len__(self):
        return len(self.X)


training_loader = DataLoader(ToyDataset(), batch_size=100, shuffle=True)


def train():
    criterion = nn.MSELoss()
    for i in range(500):
        for x, y in training_loader:
            loss = criterion(model(x.to("cuda")), y.to("cuda"))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    print("epoch #%d: loss: %f" % (i, loss.item()))


def test():
    x = torch.tensor([[2, 4, 6, 8, 10, 12, 14, 16]]).float()
    y_hat = model(x.to("cuda"))
    print("orig: ", x, " new: ", y_hat)


def get_model(f):
    global model, optimizer
    if os.path.exists(f):
        model = torch.load(f).cuda()
        optimizer = optim.Adam(model.parameters(), lr=0.01)
    else:
        model = nn.Sequential(
            nn.Linear(8, 16),
            nn.PReLU(),
            nn.Linear(16, 8),
        ).cuda()
        optimizer = optim.Adam(model.parameters(), lr=0.01)
        train()
        torch.save(model, f)


get_model("/tmp/model.pt")
print("-------orig---------")
test()

print(model[2].state_dict())
ASP.prune_trained_model(model, optimizer)
print("-------pruned---------")
test()
print(model[2].state_dict())
train()
print("-------retrain---------")
test()
print(model[2].state_dict())
torch.save(model, "/tmp/model_sparse.pt")