Spaces:
Running
Running
File size: 4,060 Bytes
85ce65e |
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 |
"""Implementation of the hard Concrete distribution.
Originally from:
https://github.com/asappresearch/flop/blob/master/flop/hardconcrete.py
"""
import math
import torch
import torch.nn as nn
class HardConcrete(nn.Module):
"""A HarcConcrete module.
Use this module to create a mask of size N, which you can
then use to perform L0 regularization.
To obtain a mask, simply run a forward pass through the module
with no input data. The mask is sampled in training mode, and
fixed during evaluation mode, e.g.:
>>> module = HardConcrete(n_in=100)
>>> mask = module()
>>> norm = module.l0_norm()
"""
def __init__(
self,
n_in: int,
init_mean: float = 0.5,
init_std: float = 0.01,
temperature: float = 2/3, # from CoFi
stretch: float = 0.1,
eps: float = 1e-6
) -> None:
"""Initialize the HardConcrete module.
Parameters
----------
n_in : int
The number of hard concrete variables in this mask.
init_mean : float, optional
Initial drop rate for hard concrete parameter,
by default 0.5.,
init_std: float, optional
Used to initialize the hard concrete parameters,
by default 0.01.
temperature : float, optional
Temperature used to control the sharpness of the
distribution, by default 1.0
stretch : float, optional
Stretch the sampled value from [0, 1] to the interval
[-stretch, 1 + stretch], by default 0.1.
"""
super().__init__()
self.n_in = n_in
self.limit_l = -stretch
self.limit_r = 1.0 + stretch
self.log_alpha = nn.Parameter(torch.zeros(n_in))
self.beta = temperature
self.init_mean = init_mean
self.init_std = init_std
self.bias = -self.beta * math.log(-self.limit_l / self.limit_r)
self.eps = eps
self.compiled_mask = None
self.reset_parameters()
def reset_parameters(self):
"""Reset the parameters of this module."""
self.compiled_mask = None
mean = math.log(1 - self.init_mean) - math.log(self.init_mean)
self.log_alpha.data.normal_(mean, self.init_std)
def l0_norm(self) -> torch.Tensor:
"""Compute the expected L0 norm of this mask.
Returns
-------
torch.Tensor
The expected L0 norm.
"""
return (self.log_alpha + self.bias).sigmoid().sum()
def forward(self) -> torch.Tensor:
"""Sample a hard concrete mask.
Returns
-------
torch.Tensor
The sampled binary mask
"""
if self.training:
# Reset the compiled mask
self.compiled_mask = None
# Sample mask dynamically
u = self.log_alpha.new(self.n_in).uniform_(self.eps, 1 - self.eps)
s = torch.sigmoid((torch.log(u / (1 - u)) + self.log_alpha) / self.beta)
s = s * (self.limit_r - self.limit_l) + self.limit_l
mask = s.clamp(min=0., max=1.)
else:
# Compile new mask if not cached
if self.compiled_mask is None:
# Get expected sparsity
expected_num_zeros = self.n_in - self.l0_norm().item()
num_zeros = round(expected_num_zeros)
# Approximate expected value of each mask variable z;
# We use an empirically validated magic number 0.8
soft_mask = torch.sigmoid(self.log_alpha / self.beta * 0.8)
# Prune small values to set to 0
_, indices = torch.topk(soft_mask, k=num_zeros, largest=False)
soft_mask[indices] = 0.
self.compiled_mask = soft_mask
mask = self.compiled_mask
return mask
def extra_repr(self) -> str:
return str(self.n_in)
def __repr__(self) -> str:
return "{}({})".format(self.__class__.__name__, self.extra_repr())
|