Skip to content

Commit 9948dba

Browse files
ssjhvcopybara-github
ssjhv
authored and
copybara-github
committed
Prepare tensorflow_compression_ops package for TF 2.15.
Due to Keras changes in TF 2.15, it would be difficult for us to provide TFC compatibility with TF 2.15 and beyond. Instead, we will build a separate Pip package tensorflow_compression_ops which contains C++-based custom ops of TFC. PiperOrigin-RevId: 597057231 Change-Id: Iaa019febab938825f85e9a053565866192195021
1 parent 6821906 commit 9948dba

File tree

9 files changed

+497
-2
lines changed

9 files changed

+497
-2
lines changed

‎tensorflow_compression/python/ops/gen_ops.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,10 @@
3232
"entropy_encode_finalize",
3333
"entropy_encode_index",
3434
"pmf_to_quantized_cdf",
35-
"run_length_gamma_decode",
36-
"run_length_gamma_encode",
35+
"run_length_decode",
36+
"run_length_encode",
37+
"run_length_gamma_decode", # Deprecated.
38+
"run_length_gamma_encode", # Deprecated.
3739
"stochastic_round",
3840
]
3941
# pylint:enable=undefined-all-variable

‎tensorflow_compression_ops/MANIFEST.in

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
global-include LICENSE
2+
global-include *.md
3+
recursive-include tensorflow_compression_ops/cc/ *.so

‎tensorflow_compression_ops/README.md

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
1+
# TensorFlow Compression Ops
2+
3+
TensorFlow Compression Ops (TFC-ops) contains data compression ops for
4+
TensorFlow.
5+
6+
This is a subset package of TensorFlow Compression (TFC) that contains
7+
C++-implemented TensorFlow operations only. For the full TFC package, please
8+
refer to the [TFC homepage](https://github.com/tensorflow/compression/).
9+
10+
11+
## Documentation & getting help
12+
13+
Refer to [the TFC API
14+
documentation](https://www.tensorflow.org/api_docs/python/tfc) for a complete
15+
description of the functions this package implements.
16+
17+
This subset pockage implements the following functions in the API:
18+
19+
* `create_range_encoder`
20+
* `create_range_decoder`
21+
* `entropy_decode_channel`
22+
* `entropy_decode_finalize`
23+
* `entropy_decode_index`
24+
* `entropy_encode_channel`
25+
* `entropy_encode_finalize`
26+
* `entropy_encode_index`
27+
* `pmf_to_quantized_cdf`
28+
* `range_decode` (deprecated)
29+
* `range_encode` (deprecated)
30+
* `run_length_decode`
31+
* `run_length_encode`
32+
* `run_length_gamma_decode` (deprecated)
33+
* `run_length_gamma_encode` (deprecated)
34+
* `stochastic_round`
35+
36+
Please post all questions or comments on
37+
[Discussions](https://github.com/tensorflow/compression/discussions). Only file
38+
[Issues](https://github.com/tensorflow/compression/issues) for actual bugs or
39+
feature requests. On Discussions, you may get a faster answer, and you help
40+
other people find the question or answer more easily later.
41+
42+
43+
## Installation
44+
45+
***Note: Precompiled packages are currently only provided for Linux and
46+
Darwin/Mac OS.***
47+
48+
Set up an environment in which you can install precompiled binary Python
49+
packages using the `pip` command. Refer to the
50+
[TensorFlow installation instructions](https://www.tensorflow.org/install/pip)
51+
for more information on how to set up such a Python environment.
52+
53+
The current version of TensorFlow Compression requires TensorFlow 2.
54+
55+
### pip
56+
57+
To install TFC via `pip`, run the following command:
58+
59+
```bash
60+
python -m pip install tensorflow-compression-ops
61+
```
62+
63+
To test that the installation works correctly, you can run the unit tests with:
64+
65+
```bash
66+
python -m tensorflow_compression_ops.tests.all
67+
```
68+
69+
Once the command finishes, you should see a message ```OK (skipped=2)``` or
70+
similar in the last line.
71+
72+
### Colab
73+
74+
You can try out TFC live in a [Colab](https://colab.research.google.com/). The
75+
following command installs the latest version of TFC that is compatible with the
76+
installed TensorFlow version. Run it in a cell before executing your Python
77+
code:
78+
79+
```
80+
%pip install tensorflow-compression-ops~=$(pip show tensorflow | perl -p -0777 -e 's/.*Version: (\d+\.\d+).*/\1.0/sg')
81+
```
82+
83+
Note: The binary packages of TFC are tied to TF with the same minor version
84+
(e.g., TFC 2.9.1 requires TF 2.9.x), and Colab sometimes lags behind a few days
85+
in deploying the latest version of TensorFlow. As a result, using `%pip install
86+
tensorflow-compression-ops` naively might attempt to upgrade TF, which may
87+
create problems.
88+
89+
### Docker
90+
91+
To use a Docker container (e.g. on Windows), be sure to install Docker
92+
(e.g., [Docker Desktop](https://www.docker.com/products/docker-desktop)),
93+
use a [TensorFlow Docker image](https://www.tensorflow.org/install/docker),
94+
and then run the `pip install` command inside the Docker container, not on the
95+
host. For instance, you can use a command line like this:
96+
97+
```bash
98+
docker run tensorflow/tensorflow:latest bash -c \
99+
"python -m pip install tensorflow-compression-ops &&
100+
python -m tensorflow_compression_ops.tests.all"
101+
```
102+
103+
This will fetch the TensorFlow Docker image if it's not already cached, install
104+
the pip package and then run the unit tests to confirm that it works.
105+
106+
### Anaconda
107+
108+
It seems that [Anaconda](https://www.anaconda.com/distribution/) ships its own
109+
binary version of TensorFlow which is incompatible with our pip package. To
110+
solve this, always install TensorFlow via `pip` rather than `conda`. For
111+
example, this creates an Anaconda environment with CUDA libraries, and then
112+
installs TensorFlow and TensorFlow Compression Ops:
113+
114+
```bash
115+
conda create --name ENV_NAME python cudatoolkit cudnn
116+
conda activate ENV_NAME
117+
python -m pip install tensorflow-compression-ops
118+
```
119+
120+
Depending on the requirements of the `tensorflow` pip package, you may need to
121+
pin the CUDA libraries to specific versions. If you aren't using a GPU, CUDA is
122+
of course not necessary.
123+
124+
125+
## Usage
126+
127+
We recommend importing the library from your Python code as follows:
128+
129+
```python
130+
import tensorflow as tf
131+
import tensorflow_compression_ops as tfc
132+
```
133+
134+
135+
## Building pip packages
136+
137+
This section describes the necessary steps to build your own pip packages of
138+
TensorFlow Compression Ops. This may be necessary to install it on platforms for
139+
which we don't provide precompiled binaries (currently only Linux and Darwin).
140+
141+
To be compatible with the official TensorFlow pip package, the TFC pip package
142+
must be linked against a matching version of the C libraries. For this reason,
143+
to build the official Linux pip packages, we use [these Docker
144+
images](https://hub.docker.com/r/tensorflow/build) and use the same toolchain
145+
that TensorFlow uses.
146+
147+
Inside the Docker container, the following steps need to be taken:
148+
149+
1. Clone the `tensorflow/compression` repo from GitHub.
150+
2. Run `tensorflow_compression_ops/build_pip_pkg.sh` inside the cloned repo.
151+
152+
For example:
153+
154+
```bash
155+
git clone https://github.com/tensorflow/compression.git /tensorflow_compression
156+
docker run -i --rm \
157+
-v /tmp/tensorflow_compression_ops:/tmp/tensorflow_compression_ops \
158+
-v /tensorflow_compression:/tensorflow_compression \
159+
-w /tensorflow_compression \
160+
-e "BAZEL_OPT=--config=manylinux_2_17_x86_64" \
161+
tensorflow/build:latest-python3.10 \
162+
bash tensorflow_compression_ops/build_pip_pkg.sh /tmp/tensorflow_compression_ops <custom-version>
163+
```
164+
165+
For Darwin, the Docker image and specifying the Bazel config is not necessary.
166+
We just build the package like this (note that you may want to create a clean
167+
Python virtual environment to do this):
168+
169+
```bash
170+
git clone https://github.com/tensorflow/compression.git /tensorflow_compression
171+
cd /tensorflow_compression
172+
BAZEL_OPT="--macos_minimum_os=10.14" bash \
173+
tensorflow_compression_ops/build_pip_pkg.sh \
174+
/tmp/tensorflow_compression_ops <custom-version>
175+
```
176+
177+
In both cases, the wheel file is created inside `/tmp/tensorflow_compression_ops`.
178+
179+
To test the created package, first install the resulting wheel file:
180+
181+
```bash
182+
python -m pip install /tmp/tensorflow_compression_ops/tensorflow_compression_ops-*.whl
183+
```
184+
185+
Then run the unit tests (Do not run the tests in the workspace directory where
186+
the `WORKSPACE` file lives. In that case, the Python interpreter would attempt
187+
to import `tensorflow_compression_ops` packages from the source tree, rather
188+
than from the installed package system directory):
189+
190+
```bash
191+
pushd /tmp
192+
python -m tensorflow_compression_ops.tests.all
193+
popd
194+
```
195+
196+
When done, you can uninstall the pip package again:
197+
198+
```bash
199+
python -m pip uninstall tensorflow-compression-ops
200+
```
201+
202+
203+
## Citation
204+
205+
If you use this library for research purposes, please cite:
206+
207+
```
208+
@software{tfc_github,
209+
author = "Ballé, Johannes and Hwang, Sung Jin and Agustsson, Eirikur",
210+
title = "{T}ensor{F}low {C}ompression: Learned Data Compression",
211+
url = "http://github.com/tensorflow/compression",
212+
version = "2.15.0",
213+
year = "2024",
214+
}
215+
```
216+
217+
In the above BibTeX entry, names are top contributors sorted by number of
218+
commits. Please adjust version number and year according to the version that was
219+
actually used.
220+
221+
Note that this is not an officially supported Google product.

‎tensorflow_compression_ops/__init__.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Copyright 2024 Google LLC. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ==============================================================================
15+
"""Data compression ops in TensorFlow."""
16+
17+
from tensorflow.python.framework import load_library
18+
from tensorflow.python.platform import resource_loader
19+
20+
ops = load_library.load_op_library(resource_loader.get_path_to_datafile(
21+
"cc/libtensorflow_compression.so"))
22+
23+
globals().update({n: getattr(ops, n) for n in [
24+
"create_range_encoder",
25+
"create_range_decoder",
26+
"entropy_decode_channel",
27+
"entropy_decode_finalize",
28+
"entropy_decode_index",
29+
"entropy_encode_channel",
30+
"entropy_encode_finalize",
31+
"entropy_encode_index",
32+
"pmf_to_quantized_cdf",
33+
"range_decode", # Deprecated.
34+
"range_encode", # Deprecated.
35+
"run_length_decode",
36+
"run_length_encode",
37+
"run_length_gamma_decode", # Deprecated.
38+
"run_length_gamma_encode", # Deprecated.
39+
"stochastic_round",
40+
]})
41+

‎tensorflow_compression_ops/build_pip_pkg.py

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
# Copyright 2024 Google LLC. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ==============================================================================
15+
"""Setup for pip package."""
16+
17+
import atexit
18+
import glob
19+
import os
20+
import shutil
21+
import sys
22+
import tempfile
23+
import setuptools
24+
25+
# Version string should follow PEP440 rules.
26+
DEFAULT_VERSION = "0.dev0+build-from-source"
27+
28+
29+
class BinaryDistribution(setuptools.Distribution):
30+
"""This class is needed in order to create OS specific wheels."""
31+
32+
def has_ext_modules(self):
33+
return True
34+
35+
36+
def main(srcdir: str, destdir: str, version: str = ""):
37+
tempdir = tempfile.mkdtemp()
38+
atexit.register(shutil.rmtree, tempdir)
39+
40+
pkgdir = os.path.join(tempdir, "tensorflow_compression_ops")
41+
shutil.copytree(os.path.join(srcdir, "tensorflow_compression_ops"), pkgdir)
42+
shutil.copy2(os.path.join(srcdir, "MANIFEST.in"), tempdir)
43+
shutil.copy2(os.path.join(srcdir, "LICENSE"), pkgdir)
44+
shutil.copy2(os.path.join(srcdir, "README.md"), pkgdir)
45+
46+
if not os.path.exists(
47+
os.path.join(pkgdir, "cc/libtensorflow_compression.so")):
48+
raise RuntimeError("libtensorflow_compression.so not found. "
49+
"Did you 'bazel run?'")
50+
51+
with open(os.path.join(srcdir, "requirements.txt"), "r") as f:
52+
install_requires = f.readlines()
53+
54+
print("=== Building wheel")
55+
atexit.register(os.chdir, os.getcwd())
56+
os.chdir(tempdir)
57+
setuptools.setup(
58+
name="tensorflow_compression_ops",
59+
version=version or DEFAULT_VERSION,
60+
description="Data compression ops for TensorFlow",
61+
url="https://tensorflow.github.io/compression/",
62+
author="Google LLC",
63+
# Contained modules and scripts.
64+
packages=setuptools.find_packages(),
65+
install_requires=install_requires,
66+
script_args=["sdist", "bdist_wheel"],
67+
# Add in any packaged data.
68+
include_package_data=True,
69+
zip_safe=False,
70+
distclass=BinaryDistribution,
71+
# PyPI package information.
72+
classifiers=[
73+
"Development Status :: 5 - Production/Stable",
74+
"Intended Audience :: Developers",
75+
"Intended Audience :: Education",
76+
"Intended Audience :: Science/Research",
77+
"License :: OSI Approved :: Apache Software License",
78+
"Programming Language :: Python :: 3",
79+
"Topic :: Scientific/Engineering :: Mathematics",
80+
"Topic :: Software Development :: Libraries :: Python Modules",
81+
"Topic :: Software Development :: Libraries",
82+
],
83+
project_urls={
84+
"Documentation":
85+
"https://tensorflow.github.io/compression/docs/api_docs/python/tfc.html",
86+
"Discussion":
87+
"https://groups.google.com/forum/#!forum/tensorflow-compression",
88+
"Source": "https://github.com/tensorflow/compression",
89+
"Tracker": "https://github.com/tensorflow/compression/issues",
90+
},
91+
license="Apache 2.0",
92+
keywords=("compression data-compression tensorflow machine-learning "
93+
"python deep-learning deep-neural-networks neural-network ml")
94+
)
95+
96+
print("=== Copying wheel to:", destdir)
97+
os.makedirs(destdir, exist_ok=True)
98+
for path in glob.glob(os.path.join(tempdir, "dist", "*.whl")):
99+
print("Copied into:", shutil.copy(path, destdir))
100+
101+
102+
if __name__ == "__main__":
103+
main(*sys.argv[1:]) # pylint: disable=too-many-function-args
104+

0 commit comments

Comments
 (0)