TensorFlow 压缩

TensorFlow Compression (TFC) 是一个为 TensorFlow 提供的数据压缩工具库。

您可以使用该库构建自己的机器学习模型，并内置端到端优化的数据压缩功能。这有助于在仅牺牲少量模型性能的情况下，找到数据（如图像、特征、样本等）的高效存储表示。请参阅 有损数据压缩教程 或 模型压缩教程，以开始您的实践。

若想从经典数据压缩的角度深入了解，可以阅读我们的 非线性变换编码论文，或观看 @jonarchists 的 关于学习型图像压缩的演讲。而从机器学习视角介绍有损数据压缩的内容，则可参考 @yiboyang 的 综述论文。

该库包含以下内容（详情请参阅 API 文档）：

• 范围编码（又称算术编码）实现，以 C++ 编写的灵活 TF 操作形式提供。这些实现包括可选的“溢出”功能，可在范围编码的比特序列中嵌入 Elias gamma 编码，从而支持对包含所有有符号整数而非有限范围的字符集进行编码。

• 熵模型类，简化了设计率–失真优化编码的过程。在训练期间，它们充当似然模型；训练完成后，通过自动设计范围编码表并调用底层范围编码实现，将浮点张量编码为优化的比特序列。

• 其他在学习型数据压缩上下文中非常有用的 TensorFlow 函数和 Keras 层，例如用于数值计算密度函数分位数的方法、基于抖动噪声取期望值的操作、具有更灵活填充选项且支持在傅里叶域中重参数化卷积核和偏置的卷积层，以及广义除法归一化（GDN）的实现。

重要更新： 自 2024 年 2 月 1 日起，TensorFlow Compression 已进入维护模式。具体而言：

• TFC 的完整功能集已冻结。不会开发新功能，但仓库将继续接收维护性修复。

• 今后，新的 TFC 包将仅兼容 TensorFlow 2.14。这是由于 TensorFlow 2.15 携带的 Keras 版本引入了不兼容性，需要重写我们的层和熵模型类。

• 为确保现有模型仍能在 TensorFlow 2.15 及更高版本上运行，我们发布了新的包 tensorflow-compression-ops，其中仅包含 C++ 操作。这些操作将尽可能地更新以支持较新的 TensorFlow 版本。

  • 更新：由于在维护 C++ 自定义操作时面临的技术挑战，TF 2.18 将是 tensorflow-compression-ops 支持的最后一个版本。

• 在 pypi.org 上同时提供了两种选项的二进制包：TFC 和 TFC 操作。

文档与帮助

有关本包所实现的类和函数的完整说明，请参阅 API 文档。

请将所有问题或评论发布到 Discussions。仅针对实际的 bug 或功能请求提交 Issues。在 Discussions 上，您可能会更快得到回复，并且有助于其他人日后更轻松地找到相关问题或答案。

安装

注意：目前仅提供适用于 Linux 和 Darwin/Mac OS 的预编译包。若要在 Windows 上使用这些包，建议按照 WSL2 安装指南 安装 TensorFlow，或使用 TensorFlow Docker 镜像，然后再安装 Linux 版本的包。

设置一个可以使用 pip 命令安装预编译 Python 二进制包的环境。有关如何搭建此类 Python 环境的详细信息，请参阅 TensorFlow 安装指南。

当前版本的 TensorFlow Compression 需要 TensorFlow 2。对于兼容 TensorFlow 1 的版本，请参阅我们的 历史版本。

pip

要通过 pip 安装 TFC，运行以下命令：

python -m pip install tensorflow-compression

为测试安装是否正确，可以运行单元测试：

python -m tensorflow_compression.all_tests

命令执行完毕后，最后一行应显示类似 OK (skipped=29) 的消息。

Colab

您可以在 Colab 中实时试用 TFC。以下命令会安装与当前安装的 TensorFlow 版本兼容的最新 TFC 版本。在执行 Python 代码之前，在单元格中运行它：

%pip install tensorflow-compression~=$(pip show tensorflow | perl -p -0777 -e 's/.*Version: (\d+\.\d+).*/\1.0/sg')

注意：TFC 的二进制包与 TensorFlow 的次版本号绑定（例如，TFC 2.9.1 需要 TensorFlow 2.9.x），而 Colab 在部署最新版 TensorFlow 方面有时会滞后几天。因此，如果直接使用 pip install tensorflow-compression，可能会尝试升级 TensorFlow，从而引发问题。

Docker

若要在 Docker 容器中使用（例如在 Windows 上），请确保已安装 Docker（如 Docker Desktop），并使用 TensorFlow Docker 镜像。然后在 Docker 容器内运行 pip install 命令，而不是在主机上。例如，可以使用如下命令：

docker run tensorflow/tensorflow:latest bash -c \
    "python -m pip install tensorflow-compression &&
     python -m tensorflow_compression.all_tests"

此命令将拉取 TensorFlow Docker 镜像（如果尚未缓存），安装 pip 包，并运行单元测试以确认其正常工作。

Anaconda

似乎 Anaconda 自带的 TensorFlow 二进制版本与我们通过 pip 发布的包不兼容。为了解决这个问题，请始终使用 pip 而不是 conda 来安装 TensorFlow。例如，以下命令会创建一个包含 CUDA 库的 Anaconda 环境，然后安装 TensorFlow 和 TensorFlow Compression：

conda create --name ENV_NAME python cudatoolkit cudnn
conda activate ENV_NAME
python -m pip install tensorflow-compression

根据 tensorflow pip 包的具体要求，你可能需要将 CUDA 库固定到特定版本。如果你不使用 GPU，则当然不需要 CUDA。

使用方法

我们建议在 Python 代码中按如下方式导入该库：

import tensorflow as tf
import tensorflow_compression as tfc

使用预训练模型压缩图像

在 models 目录 中，你会找到一个名为 tfci.py 的 Python 脚本。下载该文件并运行：

python tfci.py -h

这将列出所有可用选项。简而言之，使用以下命令：

python tfci.py compress <model> <PNG 文件>

即可利用预训练模型对图像进行压缩，并生成一个以 .tfci 结尾的文件。执行 python tfci.py models 可查看支持的预训练模型列表。而以下命令：

python tfci.py decompress <TFCI 文件>

则可以解压缩 TFCI 文件并生成一个 PNG 文件。默认情况下，输出文件名与输入文件相同，仅添加相应的文件扩展名（现有扩展名不会被移除）。

训练自己的模型

models 目录 中包含了多个已发表的图像压缩模型实现，方便用户进行实验。需要注意的是，要复现已发表的结果，可能还需要对代码和训练数据集进行更多调优。你可以使用上述 tfci.py 脚本来访问这些公开模型。

以下说明介绍了一种对以下论文中提出的模型的重新实现：

“端到端优化的图像压缩”
J. Ballé, V. Laparra, E. P. Simoncelli
https://arxiv.org/abs/1611.01704

请注意，models 目录并不包含在 pip 包中。这些模型需要单独下载。下载 bls2017.py 文件并运行：

python bls2017.py -h

这将列出该实现的所有可用命令行选项。训练过程可以非常简单，例如：

python bls2017.py -V train

此命令将使用默认设置。请注意，除非通过 --train_glob 提供自定义训练数据集，否则将会使用 TensorFlow Datasets 下载 CLIC 数据集。

最重要的训练参数是 --lambda，它控制着模型在比特率和失真之间的权衡。每层的通道数也很重要：针对较高比特率（或等价地，较低失真）优化的模型通常需要具有更高逼近能力的变换（即更多的通道）。因此，为了优化性能，你需要确保通道数足够多（或更多）。这一点在以下文献中也有详细描述：

“用于有损图像压缩的有效非线性变换”
J. Ballé
https://arxiv.org/abs/1802.00847

如果需要，你可以使用 TensorBoard 监控训练进度。为此，在开始训练之前先在后台启动一个 TensorBoard 实例，然后在浏览器中访问你机器上的 6006 端口：

tensorboard --logdir=/tmp/train_bls2017 &

训练完成后，Python 脚本会将训练好的模型保存到由 --model_path 指定的目录中（默认为当前目录下的 bls2017），格式为 TensorFlow 的 SavedModel。之后，你可以使用该脚本进行图像的压缩和解压缩操作。两个命令必须能够访问同一个保存的模型。

python bls2017.py [options] compress original.png compressed.tfci
python bls2017.py [options] decompress compressed.tfci reconstruction.png

构建 pip 包

本节介绍了构建你自己的 TensorFlow Compression pip 包所需的步骤。这可能对于在我们尚未提供预编译二进制文件的平台上安装该库的情况是必要的（目前仅支持 Linux 和 Darwin）。

为了与官方 TensorFlow pip 包兼容，TFC 的 pip 包必须链接到匹配版本的 C 库。因此，在构建官方的 Linux 版本时，我们使用 这些 Docker 镜像，并采用与 TensorFlow 相同的工具链。

在 Docker 容器内，需要执行以下步骤：

1. 从 GitHub 克隆 tensorflow/compression 仓库。
2. 在克隆的仓库中运行 tools/build_pip_pkg.sh。

例如：

git clone https://github.com/tensorflow/compression.git /tensorflow_compression
docker run -i --rm \
    -v /tmp/tensorflow_compression:/tmp/tensorflow_compression\
    -v /tensorflow_compression:/tensorflow_compression \
    -w /tensorflow_compression \
    -e "BAZEL_OPT=--config=manylinux_2_24_x86_64" \
    tensorflow/build:latest-python3.10 \
    bash tools/build_pip_pkg.sh /tmp/tensorflow_compression <custom-version>

对于 Darwin 系统，无需使用 Docker 镜像或指定工具链。我们可以直接构建包，如下所示（建议先创建一个干净的 Python 虚拟环境来完成此操作）：

git clone https://github.com/tensorflow/compression.git /tensorflow_compression
cd /tensorflow_compression
BAZEL_OPT="--macos_minimum_os=10.14" bash \
  tools/build_pip_pkg.sh \
  /tmp/tensorflow_compression <custom-version>

在这两种情况下，wheel 文件都会在 /tmp/tensorflow_compression 目录下生成。

要测试生成的包，首先安装生成的 wheel 文件：

python -m pip install /tmp/tensorflow_compression/tensorflow_compression-*.whl

然后运行单元测试（请勿在包含 WORKSPACE 文件的工作目录中运行测试。否则，Python 解释器会尝试从源代码树中导入 tensorflow_compression 包，而不是从已安装的包系统目录中导入）：

pushd /tmp
python -m tensorflow_compression.all_tests
popd

测试完成后，你可以再次卸载该 pip 包：

python -m pip uninstall tensorflow-compression

评估

我们提供了多种图像压缩方法在不同色彩空间下、基于不同指标的评估结果。更多信息请参阅 results 子目录。

引用

如果您将本库用于研究目的，请引用以下文献：

@software{tfc_github,
  author = {Ballé, Jona 和 Hwang, Sung Jin 和 Agustsson, Eirikur},
  title = {{T}ensor{F}low {C}ompression: Learned Data Compression},
  url = {http://github.com/tensorflow/compression},
  version = {2.14.1},
  year = {2024},
}

在上述 BibTeX 条目中，作者按提交次数排序，列出了主要贡献者。请根据实际使用的版本号和年份进行相应调整。

请注意，本项目并非 Google 官方支持的产品。

TensorFlow Compression 快速上手指南

TensorFlow Compression (TFC) 是一个用于构建端到端优化数据压缩模型的 TensorFlow 库。它可以帮助你在仅牺牲少量模型性能的前提下，找到数据（如图像、特征等）的高效存储表示。

重要提示：自 2024 年 2 月 1 日起，TFC 进入维护模式。功能已冻结，不再开发新特性。新版包仅支持 TensorFlow 2.14。若需使用 TF 2.15+，请单独安装 tensorflow-compression-ops（仅包含 C++ 算子），但该组件最高仅支持到 TF 2.18。

环境准备

• 操作系统：预编译包目前仅支持 Linux 和 macOS (Darwin)。
  • Windows 用户建议：使用 WSL2 或 TensorFlow Docker 镜像来运行 Linux 版本。
• Python 环境：需要配置好支持 pip 安装的 Python 环境。
• 核心依赖：
  • TensorFlow 2.14（推荐，完整功能支持）。
  • 若使用 TF 2.15 - 2.18，仅支持底层算子功能。
• 注意：如果你使用 Anaconda，请务必通过 pip 安装 TensorFlow，不要使用 conda install tensorflow，以避免二进制兼容性问题。

安装步骤

方法一：使用 pip 安装（推荐）

确保已安装 compatible 版本的 TensorFlow（建议 2.14），然后运行：

python -m pip install tensorflow-compression

验证安装： 运行单元测试确认安装成功：

python -m tensorflow_compression.all_tests

若最后一行显示 OK (skipped=...) 或类似信息，则表示安装成功。

方法二：在 Google Colab 中使用

在 Colab 单元格中运行以下命令，它会自动匹配当前环境的 TensorFlow 版本安装对应的 TFC：

%pip install tensorflow-compression~=$(pip show tensorflow | perl -p -0777 -e 's/.*Version: (\d+\.\d+).*/\1.0/sg')

方法三：使用 Docker

适用于 Windows 用户或需要隔离环境的场景：

docker run tensorflow/tensorflow:latest bash -c \
    "python -m pip install tensorflow-compression &&
     python -m tensorflow_compression.all_tests"

基本使用

1. 导入库

在 Python 代码中按以下方式导入：

import tensorflow as tf
import tensorflow_compression as tfc

2. 使用预训练模型压缩图像

TFC 提供了脚本 tfci.py 用于快速测试（需从 GitHub 仓库的 models 目录下载该文件）。

查看帮助与支持的模型：

python tfci.py -h
python tfci.py models

压缩图像： 将 <model> 替换为支持的模型名称（如 bmshj2018-factorized），<input.png> 替换为你的图片路径。

python tfci.py compress <model> <input.png>

这将生成一个 .tfci 结尾的压缩文件。

解压图像：

python tfci.py decompress <output.tfci>

这将还原出一张 PNG 图片。

3. 训练自定义模型

你可以基于官方提供的复现模型（如 Ballé et al. 2017）进行训练。需先从 GitHub 下载对应的脚本（例如 bls2017.py）。

开始训练： 以下命令使用默认设置和 CLIC 数据集（自动下载）进行训练：

python bls2017.py -V train

关键参数说明：

• --lambda：控制码率（bitrate）与失真（distortion）之间的权衡。
• --train_glob：指定自定义训练数据集路径。
• --model_path：指定训练完成后模型保存的路径。

监控训练进度： 启动 TensorBoard 查看训练曲线：

tensorboard --logdir=/tmp/train_bls2017 &

然后在浏览器访问 http://localhost:6006。

使用训练好的模型： 训练结束后，使用生成的 SavedModel 进行压缩和解压：

# 压缩
python bls2017.py [options] compress original.png compressed.tfci

# 解压
python bls2017.py [options] decompress compressed.tfci reconstruction.png