|
- # Copyright (c) Nanjing University, Vision Lab.
- # Last update: 2019.10.07
-
- #import tensorflow as tf
- import numpy as np
- #from tensorflow.contrib.coder.python.ops import coder_ops
- from tensorlayer.layers import Module
- import tensorlayer as tl
- import torchac
- import torch
-
- def myabs(x):
- pos = tl.ops.Maximum()(x, 0)
- neg = tl.ops.Minimum()(x, 0)
- out = pos + neg*(-1)
- return out
-
- class SymmetricConditional(Module):
- """Symmetric conditional entropy model.
- Argument:
- likelihood_bound;
- range_coder_precision;
- """
-
- def __init__(self, likelihood_bound=1e-9, range_coder_precision=16): #已改
- super(SymmetricConditional, self).__init__()
- self._likelihood_bound = float(likelihood_bound)
- self._range_coder_precision = int(range_coder_precision)
-
- def _standardized_cumulative(self, inputs, loc, scale): #已改
- """
- Laplace cumulative densities function.
- """
-
- mask_r = tl.ops.greater(inputs,loc)
- mask_l = tl.ops.less_equal(inputs,loc) #返回元素的真值(x <= y) bool类型的张量
- c_l = 1.0/2.0 * tl.ops.exp(-myabs(inputs - loc) / scale)
- c_r = 1.0 - 1.0/2.0 * tl.ops.exp(-myabs(inputs - loc) / scale)
- c = c_l*tl.ops.cast(mask_l,tl.float32) + c_r*tl.ops.cast(mask_r,tl.float32)
-
- return c
-
- def _likelihood(self, inputs, loc, scale): #已改
- """ Estimate the likelihoods conditioned on assumed distribution.
-
- Arguments:
- inputs;(quantized values); loc; scale;
- Return:
- likelihood.
- """
-
- # CDF.
- upper = inputs + .5
- lower = inputs - .5
-
- sign = tl.ops.Sign()(upper + lower - loc)
- upper = - sign * (upper - loc) + loc
- lower = - sign * (lower - loc) + loc
-
- cdf_upper = self._standardized_cumulative(upper, loc, scale) #用拉普拉斯分布算累积概率
- cdf_lower = self._standardized_cumulative(lower, loc, scale)
-
- likelihood = myabs(cdf_upper - cdf_lower)
-
- return likelihood
-
- def _quantize(self, inputs, mode): #已改
- """Add noise or quantize."""
- half = tl.ops.constant(0.5)
-
- if mode == "noise":
- noise = tl.initializers.random_uniform(minval=-half, maxval=half)(shape=inputs.shape, dtype=tl.float32)
- return tl.ops.add_n([inputs, noise]) #张量逐元素相加
-
- if mode == "symbols":
- outputs = tl.ops.round(inputs)
- return outputs
-
- def forward(self, inputs, loc, scale, training=True): #已改
- """Pass a tensor through the bottleneck.
-
- Arguments:
- input tensor, loc, scale.
-
- Returns:
- output quantized tensor.
- likelihoods.
- """
-
- # check.
- #inputs = tf.convert_to_tensor(inputs, dtype=self.dtype)
- #loc = tf.convert_to_tensor(loc, dtype=self.dtype)
- #scale = tf.convert_to_tensor(scale, dtype=self.dtype)
- # quantize.
- outputs = self._quantize(inputs, "noise" if training else "symbols")
- # likelihood.
- likelihood = self._likelihood(outputs, loc, scale)
- likelihood_bound = tl.ops.constant(self._likelihood_bound)
- likelihood = tl.ops.Maximum()(likelihood, likelihood_bound)
-
- return outputs, likelihood
-
- def _pmf_to_cdf(self, pmf): #已改
- cdf = tl.ops.cumsum(pmf, axis=-1) #每一列都是前面列的累加和
- spatial_dimensions = pmf.shape[:-1] + (1,)
- zeros = tl.ops.zeros(spatial_dimensions, dtype=tl.float32)
- cdf_with_0 = tl.ops.concat([zeros, cdf], axis=-1)
- cdf_with_0 = tl.ops.Minimum()(cdf_with_0, 1.0)
-
- return cdf_with_0
-
- def _get_cdf(self, loc, scale, min_v, max_v, datashape): #已改
- """Get quantized cdf for compress/decompress.
-
- Arguments:
- inputs: integer tensor min_v, max_v.
- float32 tensor loc, scale. [-1, channels]
- Return:
- cdf with shape [-1, channels, symbols]
- """
-
- # shape of cdf shound be # [-1, N]
- a = tl.ops.reshape(tl.ops.range(min_v, max_v+1), [1, max_v-min_v+1])
- channels = datashape[-1]
- a = tl.ops.tile(a, [datashape[0]*datashape[1]*datashape[2]*datashape[3]*datashape[4], 1]) #复制很多份
- a = tl.ops.cast(a, tl.float32) # [-1, N]
-
- loc = tl.ops.expand_dims(loc, -1) #[-1, 1]
- scale = tl.ops.expand_dims(scale, -1) #[-1, 1]
-
- likelihood = self._likelihood(a, loc, scale)
- likelihood_bound = tl.ops.constant(self._likelihood_bound)
- likelihood = tl.ops.Maximum()(likelihood, likelihood_bound)
- pmf = likelihood
- cdf = self._pmf_to_cdf(pmf) #[-1,N+1]
-
- #cdf = 0 #######coder_ops.pmf_to_quantized_cdf(pmf, precision=self._range_coder_precision)# [-1, C, N]
- return cdf
-
- def compress(self, inputs, loc, scale): #已改
- """Compress inputs and store their binary representations into strings.
-
- Arguments:
- inputs: `Tensor` with values to be compressed. Must have shape
- [**batch size**, length, width, height, channels]
- locs & scales: same shape like inputs.
- Returns:
- compressed: String `Tensor` vector containing the compressed
- representation of each batch element of `inputs`.
- """
-
- #inputs = tf.convert_to_tensor(inputs, dtype=self.dtype)
- #loc = tf.convert_to_tensor(loc, dtype=self.dtype)
- #scale = tf.convert_to_tensor(scale, dtype=self.dtype)
-
- datashape = inputs.shape #[ 1 16 16 16 16]
- channels = datashape[-1]
-
- # reshape.
- loc = tl.ops.reshape(loc, (-1,))
- scale = tl.ops.reshape(scale, (-1,))
- inputs = tl.ops.reshape(inputs, (-1,)) #[BatchSizexHxWxD*C]
-
- # quantize.
- values = self._quantize(inputs, "symbols")
- # get cdf
- min_v = tl.ops.cast(tl.ops.floor(tl.ops.reduce_min(values)), dtype=tl.int32) #各个维度上元素的最小值,得到一个数
- max_v = tl.ops.cast(tl.ops.ceil(tl.ops.reduce_max(values)), dtype=tl.int32)
- cdf = self._get_cdf(loc, scale, min_v, max_v, datashape) # [BatchSizexHxWxD*C, N+1]在N那个维度上:[ 0 1 2 65536]
-
- # range encode.
- values = tl.ops.cast(values, dtype=tl.int32)
- values -= min_v
- values = torch.from_numpy(tl.convert_to_numpy(values)).to(torch.int16)
- cdf = torch.from_numpy(tl.convert_to_numpy(cdf))
- strings = torchac.encode_float_cdf(cdf, values, check_input_bounds=True)
- #strings = 0 ###########coder_ops.range_encode(values, cdf, precision=self._range_coder_precision)
-
- return strings, min_v, max_v
-
- def decompress(self, strings, loc, scale, min_v, max_v, datashape): #已改
- """Decompress values from their compressed string representations.
-
- Arguments:
- strings: A string `Tensor` vector containing the compressed data.
- shape: A `Tensor` vector of int32 type. Contains the shape of the tensor to be
- decompressed. [batch size, length, width, height, channels]
- loc & scale: parameters of distributions.
- min_v & max_v: minimum & maximum values.
-
- Return: outputs [BatchSize, H, W, D, C]
- """
-
- #strings = tf.convert_to_tensor(strings, dtype=tf.string)
- datashape = tl.convert_to_tensor(datashape, dtype='int32')
- #loc = tf.convert_to_tensor(loc, dtype=self.dtype)
- #scale = tf.convert_to_tensor(scale, dtype=self.dtype)
- min_v = tl.convert_to_tensor(min_v, dtype='int32')
- max_v = tl.convert_to_tensor(max_v, dtype='int32')
-
- # reshape.
- channels = datashape[-1]
- loc = tl.ops.reshape(loc, (-1,))
- scale = tl.ops.reshape(scale, (-1,))
-
- # get cdf.
- cdf = self._get_cdf(loc, scale, min_v, max_v, datashape) # [BatchSizexHxWxDxC, N+1]
- cdf = torch.from_numpy(tl.convert_to_numpy(cdf))
- values = torchac.decode_float_cdf(cdf, strings)
-
- # range decode.
- #code_shape = (tf.reduce_prod(datashape)/channels, channels)# shape=[-1, channel]
- #values = 0 #########coder_ops.range_decode(strings, code_shape, cdf, precision=self._range_coder_precision)
- #values = tf.cast(values, tf.int32)
- values = values + min_v.numpy()
- values = tl.convert_to_tensor(values, dtype=tl.float32)
- values = tl.ops.reshape(values, datashape)
-
- return values
-
- if __name__=='__main__': #测试通过
- np.random.seed(108)
- training = False
- y = np.random.randn(2, 16, 16, 16, 16).astype("float32")*10 #标准正态分布
- conditional_entropy_model = SymmetricConditional()
- loc = np.random.randn(2, 16, 16, 16, 16).astype("float32")
- scale = np.random.rand(2, 16, 16, 16, 16).astype("float32") #服从“0~1”均匀分布
- y_gpu = tl.convert_to_tensor(y,tl.float32)
- scale = tl.ops.constant(scale)
- loc = tl.ops.constant(loc)
- scale = tl.ops.abs(scale)
- scale = tl.ops.maximum(scale, 1e-9)
- y_tilde, likelihoods = conditional_entropy_model(y, loc, scale, training)
- print("y_tilde.shape:",y_tilde.shape)
- print("likelihoods.shape:",likelihoods.shape)
- print("y_tilde[0,0,0,0]:",y_tilde[0,0,0,0])
- print("likelihoods[0,0,0,0]:",likelihoods[0,0,0,0])
- strings, min_v, max_v = conditional_entropy_model.compress(y_tilde,loc,scale) #encode
- print("min_v:",min_v)
- print("max_v:",max_v)
- y_decoded = conditional_entropy_model.decompress(strings, loc, scale, min_v, max_v, y.shape)
- compare = tl.ops.equal(tl.ops.cast(y_tilde, dtype=tl.int32),tl.ops.cast(y_decoded, dtype=tl.int32))
- compare = tl.ops.cast(compare,tl.float32)
- print("compare=False:",tl.ops.where(compare<0.1),len(tl.ops.where(compare<0.1)))
- print("y_decoded[0,0,0,0]:",y_decoded[0,0,0,0])
|