codingutf8 2020present allenai s university of illinois urbanachampaign intel nervana systems and 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 binarizers take a real value matrix as input and produce a binary values in 0 1 mask of the same shape thresholdd binarizer computes a binary mask m from a real value matrix s such that mi j 1 if and only if si j tau where tau is a real value threshold implementation is inspired from https github comarunmallyapiggyback piggyback adapting a single network to multiple tasks by learning to mask weights arun mallya dillon davis svetlana lazebnik args inputs torch floattensor the input matrix from which the binarizer computes the binary mask threshold float the threshold value in r sigmoid bool if set to true we apply the sigmoid function to the inputs matrix before comparing to threshold in this case threshold should be a value between 0 and 1 returns mask torch floattensor binary matrix of the same size as inputs acting as a mask 1 the associated weight is retained 0 the associated weight is pruned we limit the pruning so that at least 0 5 half a percent of the weights are remaining topk binarizer computes a binary mask m from a real value matrix s such that mi j 1 if and only if si j is among the k highest values of s implementation is inspired from https github comallenaihiddennetworks what s hidden in a randomly weighted neural network vivek ramanujan mitchell wortsman aniruddha kembhavi ali farhadi mohammad rastegari args inputs torch floattensor the input matrix from which the binarizer computes the binary mask threshold float the percentage of weights to keep the rest is pruned threshold is a float between 0 and 1 returns mask torch floattensor binary matrix of the same size as inputs acting as a mask 1 the associated weight is retained 0 the associated weight is pruned get the subnetwork by sorting the inputs and using the top threshold flatout and mask access the same memory magnitude binarizer computes a binary mask m from a real value matrix s such that mi j 1 if and only if si j is among the k highest values of s absolute value implementation is inspired from https github comnervanasystemsdistillerblob2291fdcc2ea642a98d4e20629acb5a9e2e04b4e6distillerpruningautomatedgradualpruner pyl24 args inputs torch floattensor the input matrix from which the binarizer computes the binary mask this input marix is typically the weight matrix threshold float the percentage of weights to keep the rest is pruned threshold is a float between 0 and 1 returns mask torch floattensor binary matrix of the same size as inputs acting as a mask 1 the associated weight is retained 0 the associated weight is pruned get the subnetwork by sorting the inputs and using the top threshold flatout and mask access the same memory coding utf 8 2020 present allenai s university of illinois urbana champaign intel nervana systems and 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 binarizers take a real value matrix as input and produce a binary values in 0 1 mask of the same shape thresholdd binarizer computes a binary mask m from a real value matrix s such that m_ i j 1 if and only if s_ i j tau where tau is a real value threshold implementation is inspired from https github com arunmallya piggyback piggyback adapting a single network to multiple tasks by learning to mask weights arun mallya dillon davis svetlana lazebnik args inputs torch floattensor the input matrix from which the binarizer computes the binary mask threshold float the threshold value in r sigmoid bool if set to true we apply the sigmoid function to the inputs matrix before comparing to threshold in this case threshold should be a value between 0 and 1 returns mask torch floattensor binary matrix of the same size as inputs acting as a mask 1 the associated weight is retained 0 the associated weight is pruned we limit the pruning so that at least 0 5 half a percent of the weights are remaining top k binarizer computes a binary mask m from a real value matrix s such that m_ i j 1 if and only if s_ i j is among the k highest values of s implementation is inspired from https github com allenai hidden networks what s hidden in a randomly weighted neural network vivek ramanujan mitchell wortsman aniruddha kembhavi ali farhadi mohammad rastegari args inputs torch floattensor the input matrix from which the binarizer computes the binary mask threshold float the percentage of weights to keep the rest is pruned threshold is a float between 0 and 1 returns mask torch floattensor binary matrix of the same size as inputs acting as a mask 1 the associated weight is retained 0 the associated weight is pruned get the subnetwork by sorting the inputs and using the top threshold flat_out and mask access the same memory magnitude binarizer computes a binary mask m from a real value matrix s such that m_ i j 1 if and only if s_ i j is among the k highest values of s absolute value implementation is inspired from https github com nervanasystems distiller blob 2291fdcc2ea642a98d4e20629acb5a9e2e04b4e6 distiller pruning automated_gradual_pruner py l24 args inputs torch floattensor the input matrix from which the binarizer computes the binary mask this input marix is typically the weight matrix threshold float the percentage of weights to keep the rest is pruned threshold is a float between 0 and 1 returns mask torch floattensor binary matrix of the same size as inputs acting as a mask 1 the associated weight is retained 0 the associated weight is pruned get the subnetwork by sorting the inputs and using the top threshold flat_out and mask access the same memory
import torch
from torch import autograd


class ThresholdBinarizer(autograd.Function):
    

    @staticmethod
    def forward(ctx, inputs: torch.tensor, threshold: float, sigmoid: bool):
        
        nb_elems = inputs.numel()
        nb_min = int(0.005 * nb_elems) + 1
        if sigmoid:
            mask = (torch.sigmoid(inputs) > threshold).type(inputs.type())
        else:
            mask = (inputs > threshold).type(inputs.type())
        if mask.sum() < nb_min:
            
            k_threshold = inputs.flatten().kthvalue(max(nb_elems - nb_min, 1)).values
            mask = (inputs > k_threshold).type(inputs.type())
        return mask

    @staticmethod
    def backward(ctx, gradOutput):
        return gradOutput, None, None


class TopKBinarizer(autograd.Function):
    

    @staticmethod
    def forward(ctx, inputs: torch.tensor, threshold: float):
        
        
        mask = inputs.clone()
        _, idx = inputs.flatten().sort(descending=True)
        j = int(threshold * inputs.numel())

        
        flat_out = mask.flatten()
        flat_out[idx[j:]] = 0
        flat_out[idx[:j]] = 1
        return mask

    @staticmethod
    def backward(ctx, gradOutput):
        return gradOutput, None


class MagnitudeBinarizer(object):
    

    @staticmethod
    def apply(inputs: torch.tensor, threshold: float):
        
        
        mask = inputs.clone()
        _, idx = inputs.abs().flatten().sort(descending=True)
        j = int(threshold * inputs.numel())

        
        flat_out = mask.flatten()
        flat_out[idx[j:]] = 0
        flat_out[idx[:j]] = 1
        return mask