CompressAI Documentation

A LatentCodec is an abstraction used to compress a latent space using various entropy modeling techniques. It functions as a miniature version of a CompressionModel, implementing the core methods forward, compress, and decompress.

By composing different LatentCodec subclasses, you can build complex entropy models (e.g., combining a side-information branch for 'z' with a conditional branch for 'y'). This modularity allows for easy swapping between different modeling strategies like factorized, hyperprior, raster-scan autoregressive, checkerboard, or channel conditional groups.

class LatentCodec(nn.Module):
    def forward(self, y: Tensor, *args, **kwargs) -> Dict[str, Any]:
        raise NotImplementedError

    def compress(self, y: Tensor, *args, **kwargs) -> Dict[str, Any]:
        raise NotImplementedError

    def decompress(
        self, strings: List[List[bytes]], shape: Any, *args, **kwargs
    ) -> Dict[str, Any]:
        raise NotImplementedError

Training a compression model requires balancing reconstruction quality (distortion) against bit-rate (rate), typically using a loss term $\mathcal{L} = \mathcal{D} + \lambda * \mathcal{R}$.

1. Rate-Distortion Loss:

   • Distortion ($\mathcal{D}$): Usually Mean Squared Error (F.mse_loss) between input $x$ and reconstruction $\hat{x}$.
   • Rate ($\mathcal{R}$): Calculated using the log-likelihoods from the entropy bottleneck. The bits-per-pixel (bpp) loss is calculated as: torch.log(y_likelihoods).sum() / (-math.log(2) * num_pixels).

2. Auxiliary Loss: The entropy bottleneck parameters must be trained to minimize the density model evaluation. This is accessed via net.entropy_bottleneck.loss() or, if using CompressionModel, via the net.aux_loss() method.

import math
import torch.nn as nn
import torch.nn.functional as F

x = torch.rand(1, 3, 64, 64)
net = Network()
x_hat, y_likelihoods = net(x)

# bitrate of the quantized latent
N, _, H, W = x.size()
num_pixels = N * H * W
bpp_loss = torch.log(y_likelihoods).sum() / (-math.log(2) * num_pixels)

# mean square error
mse_loss = F.mse_loss(x, x_hat)

# final loss term
loss = mse_loss + lmbda * bpp_loss

CompressAI uses a range Asymmetric Numeral Systems (ANS) entropy coder by default.

• To see a list of implemented entropy coders, use compressai.available_entropy_coders().
• To change the default entropy coder, use compressai.set_entropy_coder().

Manual Compression/Decompression Workflow

To manually handle the bit-stream via the entropy bottleneck:

1. Compress: Encode the image tensor, then compress the resulting latent representation into bit-strings.
2. Decompress: Decompress the bit-strings back into a latent representation (providing the original shape), then decode the latent representation back into an image tensor.

# 1. Compress an image tensor to a bit-stream
x = torch.rand(1, 3, 64, 64)
y = net.encode(x)
strings = net.entropy_bottleneck.compress(y)

# 2. Decompress a bit-stream to an image tensor
shape = y.size()[2:]
y_hat = net.entropy_bottleneck.decompress(strings, shape)
x_hat = net.decode(y_hat)

When inheriting from compressai.models.CompressionModel, you can simplify the training loop by separating the compression network parameters from the entropy bottleneck (quantiles) parameters. This allows you to use two different optimizers.

To identify the parameter groups:

• Compression parameters: All parameters except those ending in .quantiles.
• Auxiliary parameters: Parameters ending in .quantiles.

Example training loop structure:

1. Zero both optimizers.
2. Forward pass to get $\hat{x}$ and y_likelihoods.
3. Compute and backpropagate the rate-distortion loss using the main optimizer.
4. Compute and backpropagate the auxiliary loss using the auxiliary optimizer.

from compressai.models import CompressionModel
from compressai.models.utils import conv, deconv

class Network(CompressionModel):
    def __init__(self, N=128):
        super().__init__()
        self.encode = nn.Sequential(
            conv(3, N),
            GDN(N),
            conv(N, N),
            GDN(N),
            conv(N, N),
        )

        self.decode = nn.Sequential(
            deconv(N, N),
            GDN(N, inverse=True),
            deconv(N, N),
            GDN(N, inverse=True),
            deconv(N, 3),
        )

    def forward(self, x):
        y = self.encode(x)
        y_hat, y_likelihoods = self.entropy_bottleneck(y)
        x_hat = self.decode(y_hat)
        return x_hat, y_likelihoods

# Optimizer setup
parameters = set(p for n, p in net.named_parameters() if not n.endswith(".quantiles"))
aux_parameters = set(p for n, p in net.named_parameters() if n.endswith(".quantiles"))
optimizer = optim.Adam(parameters, lr=1e-4)
aux_optimizer = optim.Adam(aux_parameters, lr=1e-3)

# Training loop
x = torch.rand(1, 3, 64, 64)
for i in range(10):
    optimizer.zero_grad()
    aux_optimizer.zero_grad()

    x_hat, y_likelihoods = net(x)
    # ... compute loss ...
    loss.backward()
    optimizer.step()

    aux_loss = net.aux_loss()
    aux_loss.backward()
    aux_optimizer.step()