| import math |
|
|
| import torch |
| import numpy as np |
| from torch import nn |
| import torch.nn.functional as F |
|
|
| from ..layers.cuda_inference import build_index_dec, build_index_enc, process_with_mask |
|
|
|
|
| class EntropyCoder(): |
| def __init__(self): |
| super().__init__() |
|
|
| from MLCodec_extensions_cpp import RansEncoder, RansDecoder |
| self.encoder = RansEncoder() |
| self.decoder = RansDecoder() |
|
|
| @staticmethod |
| def pmf_to_quantized_cdf(pmf, precision=16): |
| from MLCodec_extensions_cpp import pmf_to_quantized_cdf as _pmf_to_cdf |
| cdf = _pmf_to_cdf(pmf.tolist(), precision) |
| cdf = torch.IntTensor(cdf) |
| return cdf |
|
|
| @staticmethod |
| def pmf_to_cdf(pmf, tail_mass, pmf_length, max_length): |
| entropy_coder_precision = 16 |
| cdf = torch.zeros((len(pmf_length), max_length + 2), dtype=torch.int32) |
| for i, p in enumerate(pmf): |
| prob = torch.cat((p[: pmf_length[i]], tail_mass[i]), dim=0) |
| _cdf = EntropyCoder.pmf_to_quantized_cdf(prob, entropy_coder_precision) |
| cdf[i, : _cdf.size(0)] = _cdf |
| return cdf |
|
|
| def reset(self): |
| self.encoder.reset() |
|
|
| def add_cdf(self, cdf, cdf_length, offset): |
| enc_cdf_idx = self.encoder.add_cdf(cdf, cdf_length, offset) |
| dec_cdf_idx = self.decoder.add_cdf(cdf, cdf_length, offset) |
| assert enc_cdf_idx == dec_cdf_idx |
| return enc_cdf_idx |
|
|
| def encode_y(self, symbols, cdf_group_index): |
| |
| assert symbols.dtype == torch.int16 |
| self.encoder.encode_y(symbols.cpu().numpy(), cdf_group_index) |
|
|
| def encode_z(self, symbols, cdf_group_index, start_offset, per_channel_size): |
| self.encoder.encode_z(symbols.to(torch.int8).cpu().numpy(), |
| cdf_group_index, start_offset, per_channel_size) |
|
|
| def flush(self): |
| self.encoder.flush() |
|
|
| def get_encoded_stream(self): |
| return self.encoder.get_encoded_stream().tobytes() |
|
|
| def set_stream(self, stream): |
| self.decoder.set_stream((np.frombuffer(stream, dtype=np.uint8))) |
|
|
| def decode_y(self, indexes, cdf_group_index): |
| self.decoder.decode_y(indexes.to(torch.uint8).cpu().numpy(), cdf_group_index) |
|
|
| def decode_and_get_y(self, indexes, cdf_group_index, device, dtype): |
| rv = self.decoder.decode_and_get_y(indexes.to(torch.uint8).cpu().numpy(), cdf_group_index) |
| rv = torch.as_tensor(rv) |
| return rv.to(device).to(dtype) |
|
|
| def decode_z(self, total_size, cdf_group_index, start_offset, per_channel_size): |
| self.decoder.decode_z(total_size, cdf_group_index, start_offset, per_channel_size) |
|
|
| def get_decoded_tensor(self, device, dtype, non_blocking=False): |
| rv = self.decoder.get_decoded_tensor() |
| rv = torch.as_tensor(rv) |
| return rv.to(device, non_blocking=non_blocking).to(dtype) |
|
|
| def set_use_two_entropy_coders(self, use_two_entropy_coders): |
| self.encoder.set_use_two_encoders(use_two_entropy_coders) |
| self.decoder.set_use_two_decoders(use_two_entropy_coders) |
|
|
|
|
| class Bitparm(nn.Module): |
| def __init__(self, qp_num, channel, final=False): |
| super().__init__() |
| self.final = final |
| self.h = nn.Parameter(torch.nn.init.normal_( |
| torch.empty([qp_num, channel, 1, 1]), 0, 0.01)) |
| self.b = nn.Parameter(torch.nn.init.normal_( |
| torch.empty([qp_num, channel, 1, 1]), 0, 0.01)) |
| if not final: |
| self.a = nn.Parameter(torch.nn.init.normal_( |
| torch.empty([qp_num, channel, 1, 1]), 0, 0.01)) |
| else: |
| self.a = None |
|
|
| def forward(self, x, index): |
| h = torch.index_select(self.h, 0, index) |
| b = torch.index_select(self.b, 0, index) |
| x = x * F.softplus(h) + b |
| if self.final: |
| return x |
|
|
| a = torch.index_select(self.a, 0, index) |
| return x + torch.tanh(x) * torch.tanh(a) |
|
|
|
|
| class AEHelper(): |
| def __init__(self): |
| super().__init__() |
| self.entropy_coder = None |
| self.cdf_group_index = None |
| self._offset = None |
| self._quantized_cdf = None |
| self._cdf_length = None |
|
|
| def set_cdf_info(self, quantized_cdf, cdf_length, offset): |
| self._quantized_cdf = quantized_cdf.cpu().numpy() |
| self._cdf_length = cdf_length.reshape(-1).int().cpu().numpy() |
| self._offset = offset.reshape(-1).int().cpu().numpy() |
|
|
| def get_cdf_info(self): |
| return self._quantized_cdf, \ |
| self._cdf_length, \ |
| self._offset |
|
|
|
|
| class BitEstimator(AEHelper, nn.Module): |
| def __init__(self, qp_num, channel): |
| super().__init__() |
| self.f1 = Bitparm(qp_num, channel) |
| self.f2 = Bitparm(qp_num, channel) |
| self.f3 = Bitparm(qp_num, channel) |
| self.f4 = Bitparm(qp_num, channel, True) |
| self.qp_num = qp_num |
| self.channel = channel |
|
|
| def forward(self, x, index): |
| return self.get_cdf(x, index) |
|
|
| def get_logits_cdf(self, x, index): |
| x = self.f1(x, index) |
| x = self.f2(x, index) |
| x = self.f3(x, index) |
| x = self.f4(x, index) |
| return x |
|
|
| def get_cdf(self, x, index): |
| return torch.sigmoid(self.get_logits_cdf(x, index)) |
|
|
| def update(self, entropy_coder): |
| self.entropy_coder = entropy_coder |
|
|
| with torch.no_grad(): |
| device = next(self.parameters()).device |
| medians = torch.zeros((self.qp_num, self.channel, 1, 1), device=device) |
| index = torch.arange(self.qp_num, device=device, dtype=torch.int32) |
|
|
| minima = medians + 8 |
| for i in range(8, 1, -1): |
| samples = torch.zeros_like(medians) - i |
| probs = self.forward(samples, index) |
| minima = torch.where(probs < torch.zeros_like(medians) + 0.0001, |
| torch.zeros_like(medians) + i, minima) |
|
|
| maxima = medians + 8 |
| for i in range(8, 1, -1): |
| samples = torch.zeros_like(medians) + i |
| probs = self.forward(samples, index) |
| maxima = torch.where(probs > torch.zeros_like(medians) + 0.9999, |
| torch.zeros_like(medians) + i, maxima) |
|
|
| minima = minima.int() |
| maxima = maxima.int() |
|
|
| offset = -minima |
|
|
| pmf_start = medians - minima |
| pmf_length = maxima + minima + 1 |
|
|
| max_length = pmf_length.max() |
| device = pmf_start.device |
| samples = torch.arange(max_length, device=device) |
|
|
| samples = samples[None, None, None, :] + pmf_start |
|
|
| half = float(0.5) |
|
|
| lower = self.forward(samples - half, index) |
| upper = self.forward(samples + half, index) |
| pmf = upper - lower |
|
|
| pmf = pmf[:, :, 0, :] |
| upper = self.forward(maxima.to(torch.float32), index) |
| tail_mass = lower[:, :, 0, :1] + (1.0 - upper[:, :, 0, -1:]) |
|
|
| pmf = pmf.reshape([-1, max_length]) |
| tail_mass = tail_mass.reshape([-1, 1]) |
| pmf_length = pmf_length.reshape([-1]) |
| offset = offset.reshape([-1]) |
| quantized_cdf = EntropyCoder.pmf_to_cdf(pmf, tail_mass, pmf_length, max_length) |
| cdf_length = pmf_length + 2 |
| self.set_cdf_info(quantized_cdf, cdf_length, offset) |
| self.cdf_group_index = self.entropy_coder.add_cdf(*self.get_cdf_info()) |
|
|
| def build_indexes(self, size, qp): |
| B, C, H, W = size |
| indexes = torch.arange(C, dtype=torch.int).view(1, -1, 1, 1) + qp * self.channel |
| return indexes.repeat(B, 1, H, W) |
|
|
| def encode_z(self, x, qp): |
| _, _, H, W = x.size() |
| return self.entropy_coder.encode_z(x.reshape(-1), self.cdf_group_index, qp * self.channel, |
| H * W) |
|
|
| def decode_z(self, size, qp): |
| self.entropy_coder.decode_z(self.channel * size[0] * size[1], self.cdf_group_index, |
| qp * self.channel, size[0] * size[1]) |
|
|
| def get_z(self, size, device, dtype): |
| output_size = (1, self.channel, size[0], size[1]) |
| val = self.entropy_coder.get_decoded_tensor(device, dtype, non_blocking=True) |
| return val.reshape(output_size) |
|
|
|
|
| class GaussianEncoder(AEHelper): |
| def __init__(self): |
| super().__init__() |
| self.scale_min = 0.11 |
| self.scale_max = 16.0 |
| self.scale_level = 128 |
| self.scale_table = self.get_scale_table(self.scale_min, self.scale_max, self.scale_level) |
|
|
| self.log_scale_min = math.log(self.scale_min) |
| self.log_scale_max = math.log(self.scale_max) |
| self.log_scale_step = (self.log_scale_max - self.log_scale_min) / (self.scale_level - 1) |
| self.log_step_recip = 1. / self.log_scale_step |
|
|
| self.force_zero_thres = None |
| self.decode_index_cache = {} |
| self.decode_zeros_cache = {} |
|
|
| @staticmethod |
| def get_scale_table(min_val, max_val, levels): |
| return torch.exp(torch.linspace(math.log(min_val), math.log(max_val), levels)) |
|
|
| def update(self, entropy_coder, force_zero_thres=None): |
| self.entropy_coder = entropy_coder |
| self.force_zero_thres = force_zero_thres |
|
|
| pmf_center = torch.zeros_like(self.scale_table) + 8 |
| scales = torch.zeros_like(pmf_center) + self.scale_table |
| cdf_distribution = torch.distributions.normal.Normal(0., scales) |
| for i in range(8, 1, -1): |
| samples = torch.zeros_like(pmf_center) + i |
| probs = cdf_distribution.cdf(samples) |
| probs = torch.squeeze(probs) |
| pmf_center = torch.where(probs > torch.zeros_like(pmf_center) + 0.9999, |
| torch.zeros_like(pmf_center) + i, pmf_center) |
|
|
| pmf_center = pmf_center.int() |
| pmf_length = 2 * pmf_center + 1 |
| max_length = torch.max(pmf_length).item() |
|
|
| device = pmf_center.device |
| samples = torch.arange(max_length, device=device) - pmf_center[:, None] |
| samples = samples.float() |
|
|
| scales = torch.zeros_like(samples) + self.scale_table[:, None] |
| cdf_distribution = torch.distributions.normal.Normal(0., scales) |
|
|
| upper = cdf_distribution.cdf(samples + 0.5) |
| lower = cdf_distribution.cdf(samples - 0.5) |
| pmf = upper - lower |
|
|
| tail_mass = 2 * lower[:, :1] |
|
|
| quantized_cdf = torch.Tensor(len(pmf_length), max_length + 2) |
| quantized_cdf = EntropyCoder.pmf_to_cdf(pmf, tail_mass, pmf_length, max_length) |
|
|
| self.set_cdf_info(quantized_cdf, pmf_length+2, -pmf_center) |
| self.cdf_group_index = self.entropy_coder.add_cdf(*self.get_cdf_info()) |
|
|
| def process_with_mask(self, y, scales, means, mask): |
| return process_with_mask(y, scales, means, mask, self.force_zero_thres) |
|
|
| def build_indexes_decoder(self, scales): |
| scales = scales.reshape(-1) |
| indexes, skip_cond = build_index_dec(scales, self.scale_min, self.scale_max, |
| self.log_scale_min, self.log_step_recip, |
| self.force_zero_thres) |
| if self.force_zero_thres is not None: |
| indexes = indexes[skip_cond] |
| return indexes, skip_cond |
|
|
| def build_indexes_encoder(self, symbols, scales): |
| symbols = symbols.reshape(-1) |
| scales = scales.reshape(-1) |
| symbols = build_index_enc(symbols, scales, self.scale_min, self.scale_max, |
| self.log_scale_min, self.log_step_recip, self.force_zero_thres) |
| return symbols |
|
|
| def encode_y(self, x, scales): |
| symbols = self.build_indexes_encoder(x, scales) |
| return self.entropy_coder.encode_y(symbols, self.cdf_group_index) |
|
|
| def get_decode_index_cache(self, num, device): |
| if num not in self.decode_index_cache: |
| c = torch.arange(0, num, dtype=torch.int32, device=device) |
| self.decode_index_cache[num] = c |
|
|
| return self.decode_index_cache[num] |
|
|
| def get_decode_zeros_cache(self, num, device): |
| if num not in self.decode_zeros_cache: |
| c = torch.zeros(num, dtype=torch.int32, device=device) |
| self.decode_zeros_cache[num] = c |
|
|
| return self.decode_zeros_cache[num].clone() |
|
|
| def decode_and_get_y(self, scales, dtype, device): |
| indexes, skip_cond = self.build_indexes_decoder(scales) |
| self.decode_y(indexes) |
| return self.get_y(scales.shape, scales.numel(), dtype, device, skip_cond, indexes) |
|
|
| def decode_y(self, indexes): |
| self.entropy_coder.decode_y(indexes, self.cdf_group_index) |
|
|
| def get_y(self, shape, numel, dtype, device, skip_cond, indexes): |
| if len(indexes) == 0: |
| return torch.zeros(shape, dtype=dtype, device=device) |
| if skip_cond is not None: |
| curr_index = self.get_decode_index_cache(numel, device) |
| back_index = self.get_decode_zeros_cache(numel, device) |
| back_index.masked_scatter_(skip_cond, curr_index) |
| val = self.entropy_coder.get_decoded_tensor(device, dtype, non_blocking=True) |
| if skip_cond is not None: |
| y = torch.index_select(val, 0, back_index) * skip_cond |
| return y.reshape(shape) |
| return val.reshape(shape) |
|
|
