codingutf8 2020present the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license masked linear module a fully connected layer that computes an adaptive binary mask on the fly the mask binary or not is computed at each forward pass and multiplied against the weight matrix to prune a portion of the weights the pruned weight matrix is then multiplied against the inputs and if necessary the bias is added fully connected layer with on the fly adaptive mask if needed a score matrix is created to store the importance of each associated weight args infeatures int size of each input sample outfeatures int size of each output sample bias bool if set to false the layer will not learn an additive bias default true maskinit str the initialization method for the score matrix if a score matrix is needed choices constant uniform kaiming default constant maskscale float the initialization parameter for the chosen initialization method maskinit default 0 pruningmethod str method to compute the mask choices topk threshold sigmoiedthreshold magnitude l0 default topk get the mask mask weights with computed mask compute output linear layer with masked weights coding utf 8 2020 present the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license masked linear module a fully connected layer that computes an adaptive binary mask on the fly the mask binary or not is computed at each forward pass and multiplied against the weight matrix to prune a portion of the weights the pruned weight matrix is then multiplied against the inputs and if necessary the bias is added fully connected layer with on the fly adaptive mask if needed a score matrix is created to store the importance of each associated weight args in_features int size of each input sample out_features int size of each output sample bias bool if set to false the layer will not learn an additive bias default true mask_init str the initialization method for the score matrix if a score matrix is needed choices constant uniform kaiming default constant mask_scale float the initialization parameter for the chosen initialization method mask_init default 0 pruning_method str method to compute the mask choices topk threshold sigmoied_threshold magnitude l0 default topk get the mask mask weights with computed mask compute output linear layer with masked weights
import math

import torch
from torch import nn
from torch.nn import init

from .binarizer import MagnitudeBinarizer, ThresholdBinarizer, TopKBinarizer


class MaskedLinear(nn.Linear):
    

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        mask_init: str = "constant",
        mask_scale: float = 0.0,
        pruning_method: str = "topK",
    ):
        
        super(MaskedLinear, self).__init__(in_features=in_features, out_features=out_features, bias=bias)
        assert pruning_method in ["topK", "threshold", "sigmoied_threshold", "magnitude", "l0"]
        self.pruning_method = pruning_method

        if self.pruning_method in ["topK", "threshold", "sigmoied_threshold", "l0"]:
            self.mask_scale = mask_scale
            self.mask_init = mask_init
            self.mask_scores = nn.Parameter(torch.empty(self.weight.size()))
            self.init_mask()

    def init_mask(self):
        if self.mask_init == "constant":
            init.constant_(self.mask_scores, val=self.mask_scale)
        elif self.mask_init == "uniform":
            init.uniform_(self.mask_scores, a=-self.mask_scale, b=self.mask_scale)
        elif self.mask_init == "kaiming":
            init.kaiming_uniform_(self.mask_scores, a=math.sqrt(5))

    def forward(self, input: torch.tensor, threshold: float):
        
        if self.pruning_method == "topK":
            mask = TopKBinarizer.apply(self.mask_scores, threshold)
        elif self.pruning_method in ["threshold", "sigmoied_threshold"]:
            sig = "sigmoied" in self.pruning_method
            mask = ThresholdBinarizer.apply(self.mask_scores, threshold, sig)
        elif self.pruning_method == "magnitude":
            mask = MagnitudeBinarizer.apply(self.weight, threshold)
        elif self.pruning_method == "l0":
            l, r, b = -0.1, 1.1, 2 / 3
            if self.training:
                u = torch.zeros_like(self.mask_scores).uniform_().clamp(0.0001, 0.9999)
                s = torch.sigmoid((u.log() - (1 - u).log() + self.mask_scores) / b)
            else:
                s = torch.sigmoid(self.mask_scores)
            s_bar = s * (r - l) + l
            mask = s_bar.clamp(min=0.0, max=1.0)
        
        weight_thresholded = mask * self.weight
        
        return nn.functional.linear(input, weight_thresholded, self.bias)