Spaces:
Running on Zero
Running on Zero
File size: 2,904 Bytes
7405e85 | 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 | import math
import scipy.stats
import torch
import torch.distributions as td
from compressai.ops import LowerBound
from compressai.entropy_models import GaussianConditional
def _to_float32(*args):
if len(args) == 1:
return args[0].to(dtype=torch.float)
else:
return tuple([t.to(dtype=torch.float) for t in args])
@torch.autocast('cuda', enabled=False) # disable mixed precision
def _safe_log_prob_mass(distribution, x, bin_size=1, prob_clamp=1e-6):
prob_mass = distribution.cdf(x + 0.5*bin_size) - distribution.cdf(x - 0.5*bin_size)
log_prob = torch.where(
prob_mass > prob_clamp,
torch.log(prob_mass.clamp(min=1e-8)),
distribution.log_prob(x) + math.log(bin_size)
)
return log_prob
def _sanity_check_scale_table(scale_table):
assert isinstance(scale_table, torch.Tensor)
assert (scale_table.dim() == 1) and (scale_table.shape[0] >= 1) and (scale_table.min() > 0)
assert torch.equal(scale_table, torch.sort(scale_table)[0])
def gaussian_log_prob_mass(mean, scale, x, bin_size=1.0, prob_clamp=1e-6):
""" Compute log(P) of a "quantized" Normal(`mean`, `scale`) distribution evaluated at `x`,
where P = cdf(`x` + 0.5*bin_size) - cdf(`x` - 0.5*bin_size).
Args:
mean (Tensor): mean of the Gaussian
scale (Tensor): scale (standard deviation) of the Gaussian
x (Tensor): the quantized Gaussian is evaluated at `x`
bin_size (float): quantization bin size
prob_clamp (float): when prob < prob_clamp, use approximation \
to improve numerical stability.
"""
mean, scale, x = _to_float32(mean, scale, x)
assert scale.min() > 0, f'invalid scale value = {scale.min()}'
log_prob = _safe_log_prob_mass(td.Normal(mean, scale), x, bin_size, prob_clamp)
return log_prob
class DiscretizedGaussian(GaussianConditional):
""" Custom discretized gaussian.
"""
def __init__(self, scale_table=None):
"""
Args:
scale_table (torch.Tensor, optional): a 1-D tensor.
"""
super(GaussianConditional, self).__init__()
if scale_table is None:
scale_table = self._get_default_scale_table()
_sanity_check_scale_table(scale_table)
self.register_buffer("scale_table", scale_table, persistent=False)
self.tail_mass = float(1e-9)
self.lower_bound_scale = LowerBound(scale_table[0])
self.standard_gaussian = td.Normal(loc=0, scale=1)
@staticmethod
def _get_default_scale_table():
scale_table = torch.exp(torch.linspace(math.log(0.11), math.log(20.0), steps=64))
return scale_table
@staticmethod
def _standardized_quantile(quantile):
return scipy.stats.norm.ppf(quantile)
def _standardized_cumulative(self, inputs: torch.Tensor):
return self.standard_gaussian.cdf(inputs)
|