class JointAutoregressiveHierarchicalPriors(CompressionModel): r"""Joint Autoregressive Hierarchical Priors model from D. Minnen, J. Balle, G.D. Toderici: `"Joint Autoregressive and Hierarchical Priors for Learned Image Compression" <https://arxiv.org/abs/1809.02736>`_, Adv. in Neural Information Processing Systems 31 (NeurIPS 2018). Args: N (int): Number of channels M (int): Number of channels in the expansion layers (last layer of the encoder and last layer of the hyperprior decoder) """ def __init__(self, N=192, M=192, **kwargs): super().__init__(entropy_bottleneck_channels=N, **kwargs) self.g_a = nn.Sequential( conv(3, N, kernel_size=5, stride=2), GDN(N), conv(N, N, kernel_size=5, stride=2), GDN(N), conv(N, N, kernel_size=5, stride=2), GDN(N), conv(N, M, kernel_size=5, stride=2), ) self.g_s = nn.Sequential( deconv(M, N, kernel_size=5, stride=2), GDN(N, inverse=True), deconv(N, N, kernel_size=5, stride=2), GDN(N, inverse=True), deconv(N, N, kernel_size=5, stride=2), GDN(N, inverse=True), deconv(N, 3, kernel_size=5, stride=2), ) self.h_a = nn.Sequential( conv(M, N, stride=1, kernel_size=3), nn.LeakyReLU(inplace=True), conv(N, N, stride=2, kernel_size=5), nn.LeakyReLU(inplace=True), conv(N, N, stride=2, kernel_size=5), ) self.h_s = nn.Sequential( deconv(N, M, stride=2, kernel_size=5), nn.LeakyReLU(inplace=True), deconv(M, M * 3 // 2, stride=2, kernel_size=5), nn.LeakyReLU(inplace=True), conv(M * 3 // 2, M * 2, stride=1, kernel_size=3), ) self.entropy_parameters = nn.Sequential( nn.Conv2d(M * 12 // 3, M * 10 // 3, 1), nn.LeakyReLU(inplace=True), nn.Conv2d(M * 10 // 3, M * 8 // 3, 1), nn.LeakyReLU(inplace=True), nn.Conv2d(M * 8 // 3, M * 6 // 3, 1), ) self.context_prediction = MaskedConv2d(M, 2 * M, kernel_size=5, padding=2, stride=1) self.gaussian_conditional = GaussianConditional(None) self.N = int(N) self.M = int(M) def forward(self, x): y = self.g_a(x) z = self.h_a(y) z_hat, z_likelihoods = self.entropy_bottleneck(z) params = self.h_s(z_hat) y_hat = self.gaussian_conditional._quantize( # pylint: disable=protected-access y, 'noise' if self.training else 'dequantize') ctx_params = self.context_prediction(y_hat) gaussian_params = self.entropy_parameters( torch.cat((params, ctx_params), dim=1)) scales_hat, means_hat = gaussian_params.chunk(2, 1) _, y_likelihoods = self.gaussian_conditional(y, scales_hat, means=means_hat) x_hat = self.g_s(y_hat) return { 'x_hat': x_hat, 'likelihoods': { 'y': y_likelihoods, 'z': z_likelihoods }, } @classmethod def from_state_dict(cls, state_dict): """Return a new model instance from `state_dict`.""" N = state_dict['g_a.0.weight'].size(0) M = state_dict['g_a.6.weight'].size(0) net = cls(N, M) net.load_state_dict(state_dict) return net def compress(self, x): y = self.g_a(x) z = self.h_a(y) z_strings = self.entropy_bottleneck.compress(z) z_hat = self.entropy_bottleneck.decompress(z_strings, z.size()[-2:]) params = self.h_s(z_hat) s = 4 # scaling factor between z and y kernel_size = 5 # context prediction kernel size padding = (kernel_size - 1) // 2 y_height = z_hat.size(2) * s y_width = z_hat.size(3) * s y_hat = F.pad(y, (padding, padding, padding, padding)) # yapf: enable # pylint: disable=protected-access cdf = self.gaussian_conditional._quantized_cdf.tolist() cdf_lengths = self.gaussian_conditional._cdf_length.reshape( -1).int().tolist() offsets = self.gaussian_conditional._offset.reshape(-1).int().tolist() # pylint: enable=protected-access y_strings = [] for i in range(y.size(0)): encoder = BufferedRansEncoder() # Warning, this is slow... # TODO: profile the calls to the bindings... for h in range(y_height): for w in range(y_width): y_crop = y_hat[i:i + 1, :, h:h + kernel_size, w:w + kernel_size] ctx_params = self.context_prediction(y_crop) # 1x1 conv for the entropy parameters prediction network, so # we only keep the elements in the "center" ctx_p = ctx_params[i:i + 1, :, padding:padding + 1, padding:padding + 1] p = params[i:i + 1, :, h:h + 1, w:w + 1] gaussian_params = self.entropy_parameters( torch.cat((p, ctx_p), dim=1)) scales_hat, means_hat = gaussian_params.chunk(2, 1) indexes = self.gaussian_conditional.build_indexes( scales_hat) y_q = torch.round(y_crop - means_hat) y_hat[i, :, h + padding, w + padding] = (y_q + means_hat)[i, :, padding, padding] encoder.encode_with_indexes( y_q[i, :, padding, padding].int().tolist(), indexes[i, :].squeeze().int().tolist(), cdf, cdf_lengths, offsets) string = encoder.flush() y_strings.append(string) # yapf: disable return {'strings': [y_strings, z_strings], 'shape': z.size()[-2:]} def decompress(self, strings, shape): assert isinstance(strings, list) and len(strings) == 2 # FIXME: we don't respect the default entropy coder and directly call the # range ANS decoder z_hat = self.entropy_bottleneck.decompress(strings[1], shape) params = self.h_s(z_hat) s = 4 # scaling factor between z and y kernel_size = 5 # context prediction kernel size padding = (kernel_size - 1) // 2 y_height = z_hat.size(2) * s y_width = z_hat.size(3) * s # initialize y_hat to zeros, and pad it so we can directly work with # sub-tensors of size (N, C, kernel size, kernel_size) # yapf: disable y_hat = torch.zeros((z_hat.size(0), self.M, y_height + 2 * padding, y_width + 2 * padding), device=z_hat.device) decoder = RansDecoder() # pylint: disable=protected-access cdf = self.gaussian_conditional._quantized_cdf.tolist() cdf_lengths = self.gaussian_conditional._cdf_length.reshape(-1).int().tolist() offsets = self.gaussian_conditional._offset.reshape(-1).int().tolist() # Warning: this is slow due to the auto-regressive nature of the # decoding... See more recent publication where they use an # auto-regressive module on chunks of channels for faster decoding... for i, y_string in enumerate(strings[0]): decoder.set_stream(y_string) for h in range(y_height): for w in range(y_width): # only perform the 5x5 convolution on a cropped tensor # centered in (h, w) y_crop = y_hat[i:i + 1, :, h:h + kernel_size, w:w + kernel_size] # ctx_params = self.context_prediction(torch.round(y_crop)) ctx_params = self.context_prediction(y_crop) # 1x1 conv for the entropy parameters prediction network, so # we only keep the elements in the "center" ctx_p = ctx_params[i:i + 1, :, padding:padding + 1, padding:padding + 1] p = params[i:i + 1, :, h:h + 1, w:w + 1] gaussian_params = self.entropy_parameters(torch.cat((p, ctx_p), dim=1)) scales_hat, means_hat = gaussian_params.chunk(2, 1) indexes = self.gaussian_conditional.build_indexes(scales_hat) rv = decoder.decode_stream( indexes[i, :].squeeze().int().tolist(), cdf, cdf_lengths, offsets) rv = torch.Tensor(rv).reshape(1, -1, 1, 1) rv = self.gaussian_conditional._dequantize(rv, means_hat) y_hat[i, :, h + padding:h + padding + 1, w + padding:w + padding + 1] = rv y_hat = y_hat[:, :, padding:-padding, padding:-padding] # pylint: enable=protected-access # yapf: enable x_hat = self.g_s(y_hat).clamp_(0, 1) return {'x_hat': x_hat} def update(self, scale_table=None, force=False): if scale_table is None: scale_table = get_scale_table() self.gaussian_conditional.update_scale_table(scale_table, force=force) super().update(force=force) def load_state_dict(self, state_dict): # Dynamically update the entropy bottleneck buffers related to the CDFs update_registered_buffers(self.entropy_bottleneck, 'entropy_bottleneck', ['_quantized_cdf', '_offset', '_cdf_length'], state_dict) update_registered_buffers( self.gaussian_conditional, 'gaussian_conditional', ['_quantized_cdf', '_offset', '_cdf_length', 'scale_table'], state_dict) super().load_state_dict(state_dict)
class HSIC(CompressionModel): def __init__(self,N=128,M=192,K=5,**kwargs): #'cuda:0' or 'cpu' super().__init__(entropy_bottleneck_channels=N, **kwargs) # super(DSIC, self).__init__() # self.entropy_bottleneck1 = CompressionModel(entropy_bottleneck_channels=N) # self.entropy_bottleneck2 = CompressionModel(entropy_bottleneck_channels=N) self.gaussian1 = GaussianMixtureConditional(K = K) self.gaussian2 = GaussianMixtureConditional(K = K) self.N = int(N) self.M = int(M) self.K = int(K) #定义组件 self.encoder1 = Encoder1(N,M) self.encoder2 = Encoder2(N,M) self.decoder1 = Decoder1(N,M) self.decoder2 = Decoder2(N,M) # pic2 需要的组件 # #hyper # self._h_a1 = encode_hyper(N=N,M=M) # self._h_a2 = encode_hyper(N=N,M=M) # self._h_s1 = gmm_hyper_y1(N=N,M=M,K=K) # self._h_s2 = gmm_hyper_y2(N=N,M=M,K=K) ###################################################################### self.h_a1 = nn.Sequential( conv(M, N, stride=1, kernel_size=3), nn.LeakyReLU(inplace=True), conv(N, N, stride=2, kernel_size=5), nn.LeakyReLU(inplace=True), conv(N, N, stride=2, kernel_size=5), ) self.h_s1 = nn.Sequential( deconv(N, M, stride=2, kernel_size=5), nn.LeakyReLU(inplace=True), deconv(M, M * 3 // 2, stride=2, kernel_size=5), nn.LeakyReLU(inplace=True), conv(M * 3 // 2, M * 2, stride=1, kernel_size=3), ) self.entropy_parameters1 = nn.Sequential( nn.Conv2d(M * 12 // 3, M * 10 // 3, 1), nn.LeakyReLU(inplace=True), nn.Conv2d(M * 10 // 3, M * 8 // 3, 1), nn.LeakyReLU(inplace=True), nn.Conv2d(M * 8 // 3, M * 6 // 3, 1), ) self.context_prediction1 = MaskedConv2d(M, 2 * M, kernel_size=5, padding=2, stride=1) self.gaussian_conditional1 = GaussianConditional(None) self.h_a2 = nn.Sequential( conv(M, N, stride=1, kernel_size=3), nn.LeakyReLU(inplace=True), conv(N, N, stride=2, kernel_size=5), nn.LeakyReLU(inplace=True), conv(N, N, stride=2, kernel_size=5), ) self.h_s2 = nn.Sequential( deconv(N, M, stride=2, kernel_size=5), nn.LeakyReLU(inplace=True), deconv(M, M * 3 // 2, stride=2, kernel_size=5), nn.LeakyReLU(inplace=True), conv(M * 3 // 2, M * 2, stride=1, kernel_size=3), ) self.entropy_parameters2 = nn.Sequential( nn.Conv2d(M * 18 // 3, M * 10 // 3, 1), # (M * 12 // 3, M * 10 // 3, 1), nn.LeakyReLU(inplace=True), nn.Conv2d(M * 10 // 3, M * 8 // 3, 1), nn.LeakyReLU(inplace=True), nn.Conv2d(M * 8 // 3, M * 6 // 3, 1), ) self.context_prediction2 = MaskedConv2d(M, 2 * M, kernel_size=5, padding=2, stride=1) self.gaussian_conditional2 = GaussianConditional(None) def forward(self,x1,x2,h_matrix): #定义结构 y1,g1_1,g1_2,g1_3 = self.encoder1(x1) z1 = self.h_a1(y1) #print(z1.device) z1_hat,z1_likelihoods = self.entropy_bottleneck1(z1) #change: params1 = self.h_s1(z1_hat) y1_hat = self.gaussian_conditional1._quantize( # pylint: disable=protected-access y1, 'noise' if self.training else 'dequantize') ctx_params1 = self.context_prediction1(y1_hat) #用两次!! 2M gaussian_params1 = self.entropy_parameters1( torch.cat((params1, ctx_params1), dim=1)) scales_hat1, means_hat1 = gaussian_params1.chunk(2, 1) _, y1_likelihoods = self.gaussian_conditional1(y1, scales_hat1, means=means_hat1) # gmm1 = self._h_s1(z1_hat) #三要素 # y1_hat, y1_likelihoods = self.gaussian1(y1, gmm1[0],gmm1[1],gmm1[2]) # sigma x1_hat,g1_4,g1_5,g1_6 = self.decoder1(y1_hat) ############################################# #encoder x1_warp = kornia.warp_perspective(x1, h_matrix, (x1.size()[-2],x1.size()[-1])) y2 = self.encoder2(x1_warp,x2) ##end encoder # hyper for pic2 z2 = self.h_a2(y2) z2_hat, z2_likelihoods = self.entropy_bottleneck2(z2) #change params2 = self.h_s2(z2_hat) y2_hat = self.gaussian_conditional2._quantize( # pylint: disable=protected-access y2, 'noise' if self.training else 'dequantize') ctx_params2 = self.context_prediction2(y2_hat) gaussian_params2 = self.entropy_parameters2( torch.cat((params2, ctx_params2, ctx_params1), dim=1)) scales_hat2, means_hat2 = gaussian_params2.chunk(2, 1) _, y2_likelihoods = self.gaussian_conditional1(y2, scales_hat2, means=means_hat2) # gmm2 = self._h_s2(z2_hat, y1_hat) # 三要素 # y2_hat, y2_likelihoods = self.gaussian2(y2, gmm2[0], gmm2[1], gmm2[2]) # 这里也是临时,待改gmm # end hyper for pic2 ##decoder x1_hat_warp = kornia.warp_perspective(x1_hat, h_matrix, (x1_hat.size()[-2],x1_hat.size()[-1])) x2_hat = self.decoder2(y2_hat,x1_hat_warp) #end decoder # print(x1.size()) return { 'x1_hat': x1_hat, 'x2_hat': x2_hat, 'likelihoods':{ 'y1': y1_likelihoods, 'y2': y2_likelihoods, 'z1': z1_likelihoods, 'z2': z2_likelihoods, } }