decision-rules-network/sparse_linear.py at master · Joeyonng/decision-rules-network · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
import math

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.prune as prune
import torch.optim as optim

def sparse_linear(name):
    if name == 'linear':
        return Linear
    elif name == 'l0':
        return L0Linear
    elif name == 'reweight':
        return ReweightLinear
    else:
        raise ValueError(f'{name} linear type not supported.')

class Linear(nn.Linear):
    def __init__(self, in_features, out_features,  bias=True, linear=F.linear, **kwargs):
        super(Linear, self).__init__(in_features, out_features, bias=bias, **kwargs)

        self.linear = linear

    def forward(self, input):
        output = self.linear(input, self.weight, self.bias)

        return output

    def sparsity(self):
        sparsity = (self.weight == 0).float().mean().item()

        return sparsity

    def masked_weight(self):
        masked_weight = self.weight

        return masked_weight

    def regularization(self):
        regularization = 0

        return regularization

class L0Linear(nn.Linear):
    def __init__(self, in_features, out_features,  bias=True, linear=F.linear, loc_mean=0, loc_sdev=0.01,
                 beta=2 / 3, gamma=-0.1, zeta=1.1, fix_temp=True, **kwargs):
        super(L0Linear, self).__init__(in_features, out_features, bias=bias, **kwargs)

        self._size = self.weight.size()
        self.loc = nn.Parameter(torch.zeros(self._size).normal_(loc_mean, loc_sdev))
        self.temp = beta if fix_temp else nn.Parameter(torch.zeros(1).fill_(beta))
        self.register_buffer("uniform", torch.zeros(self._size))
        self.gamma = gamma
        self.zeta = zeta
        self.gamma_zeta_ratio = math.log(-gamma / zeta)
        self.linear = linear

        self.penalty = 0

    def forward(self, input):
        mask, self.penalty = self._get_mask()
        masked_weight = self.weight * mask
        output = self.linear(input, masked_weight, self.bias)

        return output

    def sparsity(self):
        sparsity = (self.masked_weight() == 0).float().mean().item()

        return sparsity

    def masked_weight(self):
        mask, _ = self._get_mask()
        masked_weight = self.weight * mask

        return masked_weight

    def regularization(self, mean=True, axis=None):
        regularization = self.penalty
        if mean:
            regularization = regularization.mean() if axis == None else regularization.mean(axis)

        return regularization

    def _get_mask(self):
        def hard_sigmoid(x):
            return torch.min(torch.max(x, torch.zeros_like(x)), torch.ones_like(x))

        if self.training:
            self.uniform.uniform_()
            u = torch.autograd.Variable(self.uniform)
            s = torch.sigmoid((torch.log(u) - torch.log(1 - u) + self.loc) / self.temp)
            s = s * (self.zeta - self.gamma) + self.gamma
            penalty = torch.sigmoid(self.loc - self.temp * self.gamma_zeta_ratio)
        else:
            s = torch.sigmoid(self.loc) * (self.zeta - self.gamma) + self.gamma
            penalty = 0

        return hard_sigmoid(s), penalty

class ReweightLinear(nn.Linear):
    def __init__(self, in_features, out_features, bias=True, linear=F.linear,
                 prune_neuron=False, prune_always=True, factor=0.1):
        super(ReweightLinear, self).__init__(in_features, out_features, bias=bias)

        self.prune_neuron = prune_neuron
        self.prune_always = prune_always
        self.factor = factor
        self.linear = linear

    def forward(self, input):
        if self.eval():
            weight = self.masked_weight()
        else:
            weight = self.masked_weight() if self.prune_always else self.weight
        out = self.linear(input, weight, self.bias)

        return out

    def sparsity(self):
        sparsity = (self.weight.abs() <= self._threshold()).float().mean().item()

        return sparsity

    def masked_weight(self):
        masked_weight = self.weight.clone()
        masked_weight[self.weight.abs() <= self._threshold()] = 0

        return masked_weight

    def regularization(self, mean=True, axis=None):
        regularization = self.weight.abs()
        if mean:
            regularization = regularization.mean() if axis == None else regularization.mean(axis)

        return regularization

    def _threshold(self):
        if self.prune_neuron:
            threshold = self.factor * self.weight.std(1).unsqueeze(1)
        else:
            threshold = self.factor * self.weight.std()

        return threshold