weight_hist.py

import torch
import plotly.graph_objects as go
import numpy as np
from torch import nn
from torch.nn import init
from scipy.stats import pearsonr
from torchvision.models.resnet import BasicBlock
from copy import deepcopy

benford = np.array([30.1, 17.6, 12.5, 9.7, 7.9, 6.7, 5.8, 5.1, 4.6]) / 100
benford_th = torch.FloatTensor(benford)


def non_bias(m, include_bn=False):
    if include_bn:
        if (
            isinstance(m, nn.Linear)
            or isinstance(m, nn.Conv2d)
            or isinstance(m, nn.BatchNorm2d)
        ):
            return m.weight
        return None
    else:
        if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
            return m.weight
        return None


def benford_r2(bin_percent):
    return pearsonr(benford, bin_percent[1:])[0]


def bincount(tensor):
    if isinstance(tensor, torch.Tensor):
        ndim = len(tensor.shape)
        if ndim == 2:
            counts = torch.zeros((tensor.shape[0], 10))
            for i in range(10):
                counts[:, i] = torch.count_nonzero(tensor == i, dim=1)
            return counts
        counts = torch.zeros(10)
        for i in range(10):
            counts[i] = torch.count_nonzero(tensor == i)
        return counts

    ndim = len(tensor.shape)
    if ndim == 2:
        counts = np.zeros(size=(tensor.shape[0], 10))
        for i in range(10):
            counts[:, i] = np.count_nonzero(tensor == i, axis=1)
        return counts
    counts = np.zeros(10)
    for i in range(10):
        counts[i] = np.count_nonzero(tensor == i)
    return counts


@torch.no_grad()
def bin_percent(tensor):
    tensor = tensor.abs() * 1e10
    tensor = tensor // 10 ** torch.log10(tensor).long()
    tensor = bincount(tensor.numpy().astype("int64"))
    return tensor / tensor.sum()


def block_bincount(net, include_bn=False):
    bins = []
    num_params = []
    total_num_params = 0
    for m in net.modules():
        # Check if leaf module
        if list(m.children()) == []:
            weight = non_bias(m, include_bn=include_bn)
            if weight is not None:
                n_param = weight.numel()
                num_params.append(n_param)
                total_num_params += n_param
                bins.append(bin_percent(weight.view(-1).detach()))

    out = torch.zeros(10)
    for b, n_param in zip(bins, num_params):
        out += (b) * (n_param / total_num_params)
    return out


def benford_r2_model(model, include_bn=False):
    bins = block_bincount(deepcopy(model).cpu(), include_bn=include_bn)
    return benford_r2(bins)


# def init_params(net, bias=False, kind="normal"):
#     """Init layer parameters."""
#     if kind == "normal":
#         kaiming = init.kaiming_normal_
#         xavier = init.normal_
#     else:
#         kaiming = init.kaiming_uniform_
#         xavier = init.uniform_

#     if bias:
#         for m in net.modules():
#             if isinstance(m, nn.Conv2d):
#                 kaiming(m.weight, mode="fan_out")
#                 if m.bias is not None:
#                     xavier(m.bias, std=1e-3)
#             elif isinstance(m, nn.BatchNorm2d):
#                 xavier(m.weight, std=1e-3)
#                 xavier(m.bias, std=1e-3)
#             elif isinstance(m, nn.Linear):
#                 kaiming(m.weight, mode="fan_out")
#                 if m.bias is not None:
#                     xavier(m.bias, std=1e-3)
#     else:
#         for m in net.modules():
#             if isinstance(m, nn.Conv2d):
#                 kaiming(m.weight, mode="fan_out")
#                 if m.bias is not None:
#                     init.constant_(m.bias, 0)
#             elif isinstance(m, nn.BatchNorm2d):
#                 init.constant_(m.weight, 1)
#                 init.constant_(m.bias, 0)
#             elif isinstance(m, nn.Linear):
#                 init.normal_(m.weight, std=1e-3)
#                 if m.bias is not None:
#                     init.constant_(m.bias, 0)
                    
arg_to_init_fn = {
    'kaiming_uniform_': init.kaiming_uniform_,
    'kaiming_normal_': init.kaiming_normal_,
    'xavier_uniform_': init.xavier_uniform_,
    'xavier_normal_': init.xavier_normal_,
    'orthogonal_': init.orthogonal_,
    'normal_': init.normal_,
    'uniform_': init.uniform_
}


def init_params(net, initializer, bias=False):
    '''Init layer parameters.'''
    init_fn = arg_to_init_fn(initializer)
    for m in net.modules():
        if isinstance(m, nn.Conv2d):
            init_fn(m.weight, mode='fan_out')
            if m.bias is not None:
                init.constant_(m.bias, 0)
        elif isinstance(m, nn.BatchNorm2d):
            init.constant_(m.weight, 1)
            init.constant_(m.bias, 0)
        elif isinstance(m, nn.Linear):
            init_fn(m.weight, std=1e-3)
            if m.bias is not None:
                init.constant_(m.bias, 0)


def get_params(m, include_bias=False):
    """Init layer parameters."""
    defs = []
    init_weights = []
    bias = 0
    if isinstance(m, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)):
        defs.append(m.weight.shape)
        init_weights.append(m.weight.detach().numpy().reshape(-1))
        if m.bias is not None and include_bias:
            bias = 1
            defs.append(m.bias.shape)
    return defs, bias, init_weights


def plot(bin_percent, title=None):
    fig = go.Figure(
        data=[
            go.Bar(x=np.arange(10)[1:], y=bin_percent[1:], name="Weights"),
            go.Scatter(x=np.arange(10)[1:], y=benford, name="Benford's Law"),
        ]
    )
    fig.update_layout(title=title)
    fig.show()


def plot_model_bar(untrained, trained, title, exclude_fc=False, fc_only=False):
    if fc_only:
        p = trained.fc.weight.view(-1).detach()
        bins_tr = bin_percent(p)

        p = untrained.fc.weight.view(-1).detach()
        bins_utr = bin_percent(p)
    else:
        p = deepcopy(trained)
        if exclude_fc:
            p.fc = None
        bins_tr = block_bincount(p)

        p = deepcopy(untrained)
        if exclude_fc:
            p.fc = None
        bins_utr = block_bincount(p)

    fig = go.Figure(
        data=[
            go.Bar(x=np.arange(10)[1:], y=bins_tr[1:], name="Trained"),
            go.Bar(x=np.arange(10)[1:], y=bins_utr[1:], name="Random"),
            go.Scatter(x=np.arange(10)[1:], y=benford, name="Benford's Law"),
        ]
    )
    fig.update_layout(title=title, barmode="group")

    print(" " * 21 + "Pearson's R v/s Benford's Law")
    print("{:20}".format("Random"), round(benford_r2(bins_utr), 4))
    print("{:20}".format("Trained"), round(benford_r2(bins_tr), 4))

    fig.show()


def plot_model_layerwise(untrained, trained, title, x, y, size, color, family, layer_names=None):
    scores1, scores2 = [], []
    flag = 0
    if layer_names is None:
        layer_names = []
        flag = 1
    i = 1
    for m1, m2 in zip(trained.children(), untrained.children()):
        if sum(p.numel() for p in m1.parameters()) == 0 or isinstance(
            m1, nn.BatchNorm2d
        ):
            continue
        score1, score2 = benford_r2(block_bincount(m1)), benford_r2(block_bincount(m2))
        scores1.append(score1)
        scores2.append(score2)
        if flag:
            layer_names.append(f"Block{i}")
        i += 1
    if flag:
        layer_names[-1] = "FC"

    fig = go.Figure(
        data=[
            go.Scatter(x=layer_names, y=scores1, name="Trained"),
            go.Scatter(x=layer_names, y=scores2, name="Random"),
        ]
    )
    fig.update_layout(legend=dict(
        yanchor="bottom",
        y=0.01,
        xanchor="left",
        x=0.01
    ))
    fig.update_layout(title=title, xaxis=dict(title=x),
     yaxis=dict(title=y), barmode="group", font=dict(
        family=family,
        size=size,
        color=color
    ))

    fig.show()


if __name__ == "__main__":
    layer1 = torch.nn.Sequential(torch.nn.Linear(784, 10), torch.nn.Linear(784, 10))
    layer2 = torch.nn.Sequential(torch.nn.Linear(784, 10), torch.nn.Linear(784, 10))

    # f = bin_percent(layer.weight.view(-1))
    # print(f)
    # print(f.shape)

    # plot(f, 'dummy title')

    # plot_model_bar(layer1, layer2, 'mlp')
    plot_model_layerwise(layer1, layer2, "layerwise")