Updates for v2.0b2 to match TF v2.4.

PiperOrigin-RevId: 350246641
Change-Id: I87f4096d2745fb280bae94b21b8465f3916d2049

Author: Johannes Ballé

README.md

@@ -48,34 +48,40 @@ packages using the `pip` command. Refer to the
 for more information on how to set up such a Python environment.
 
 The current stable version of TFC (1.3) requires TensorFlow 1.15. The current
-beta release of TFC (2.0b1) is built for TensorFlow 2.3. For versions compatible
+beta release of TFC (2.0b2) is built for TensorFlow 2.4. For versions compatible
 with TensorFlow 1.14 or earlier, see our [previous
-releases](https://github.com/tensorflow/compression/releases). The following
-instructions are for the stable release. To install the beta, replace the
-version numbers with 2.0b1 and 2.3, respectively.
+releases](https://github.com/tensorflow/compression/releases).
+
+### pip
+
+To install TF and TFC via `pip`, run the following command:
 
-You can
-install TensorFlow from any source. To install it via `pip`, run the following
-command:
-```bash
-pip install tensorflow-gpu==1.15
-```
-for GPU support, or
 ```bash
-pip install tensorflow==1.15
+pip install tensorflow-gpu==1.15 tensorflow-compression==1.3
 ```
-for CPU-only.
 
-Then, run the following command to install the tensorflow-compression pip
-package:
+for the stable release, or
+
 ```bash
-pip install tensorflow-compression==1.3
+pip install tensorflow-gpu==2.4 tensorflow-probability==0.12.1 tensorflow-compression==2.0b2
 ```
 
+for the beta release. If you don't need GPU support, you can drop the `-gpu`
+part.
+
 To test that the installation works correctly, you can run the unit tests with
+(respectively):
+
 ```bash
 python -m tensorflow_compression.python.all_test
 ```
+
+or
+
+```bash
+python -m tensorflow_compression.all_tests
+```
+
 Once the command finishes, you should see a message ```OK (skipped=12)``` or
 similar in the last line.
 
@@ -86,11 +92,21 @@ To use a Docker container (e.g. on Windows), be sure to install Docker
 use a [TensorFlow Docker image](https://www.tensorflow.org/install/docker),
 and then run the `pip install` command inside the Docker container, not on the
 host. For instance, you can use a command line like this:
+
 ```bash
 docker run tensorflow/tensorflow:1.15.0-py3 bash -c \
     "pip install tensorflow-compression==1.3 &&
      python -m tensorflow_compression.python.all_test"
 ```
+
+or (for the beta version):
+
+```bash
+docker run tensorflow/tensorflow:2.4.0 bash -c \
+    "pip install tensorflow-probability==0.12.1 tensorflow-compression==2.0b2 &&
+     python -m tensorflow_compression.all_tests"
+```
+
 This will fetch the TensorFlow Docker image if it's not already cached, install
 the pip package and then run the unit tests to confirm that it works.
 
@@ -102,6 +118,7 @@ solve this, always install TensorFlow via `pip` rather than `conda`. For
 example, this creates an Anaconda environment with Python 3.6 and CUDA
 libraries, and then installs TensorFlow and tensorflow-compression with GPU
 support:
+
 ```bash
 conda create --name ENV_NAME python=3.6 cudatoolkit=10.0 cudnn
 conda activate ENV_NAME
@@ -122,20 +139,25 @@ import tensorflow_compression as tfc
 In the
 [models directory](https://github.com/tensorflow/compression/tree/master/models),
 you'll find a python script `tfci.py`. Download the file and run:
+
 ```bash
 python tfci.py -h
 ```
 
 This will give you a list of options. Briefly, the command
+
 ```bash
 python tfci.py compress <model> <PNG file>
 ```
+
 will compress an image using a pre-trained model and write a file ending in
 `.tfci`. Execute `python tfci.py models` to give you a list of supported
 pre-trained models. The command
+
 ```bash
 python tfci.py decompress <TFCI file>
 ```
+
 will decompress a TFCI file and write a PNG file. By default, an output file
 will be named like the input file, only with the appropriate file extension
 appended (any existing extensions will not be removed).
@@ -151,13 +173,15 @@ contains an implementation of the image compression model described in:
 > https://arxiv.org/abs/1611.01704
 
 To see a list of options, download the file `bls2017.py` and run:
+
 ```bash
 python bls2017.py -h
 ```
 
 To train the model, you need to supply it with a dataset of RGB training images.
 They should be provided in PNG format. Training can be as simple as the
 following command:
+
 ```bash
 python bls2017.py --verbose train --train_glob="images/*.png"
 ```
@@ -177,13 +201,15 @@ enough (or larger). This is described in more detail in:
 If you wish, you can monitor progress with Tensorboard. To do this, create a
 Tensorboard instance in the background before starting the training, then point
 your web browser to [port 6006 on your machine](http://localhost:6006):
+
 ```bash
 tensorboard --logdir=. &
 ```
 
 When training has finished, the Python script can be used to compress and
 decompress images as follows. The same model checkpoint must be accessible to
 both commands.
+
 ```bash
 python bls2017.py [options] compress original.png compressed.tfci
 python bls2017.py [options] decompress compressed.tfci reconstruction.png
@@ -209,6 +235,7 @@ Inside a Docker container from the image, the following steps need to be taken.
 2. Run `:build_pip_pkg` inside the cloned repo.
 
 For example:
+
 ```bash
 sudo docker run -v /tmp/tensorflow_compression:/tmp/tensorflow_compression \
     tensorflow/tensorflow:nightly-custom-op-ubuntu16 bash -c \
@@ -222,6 +249,7 @@ The wheel file is created inside `/tmp/tensorflow_compression`. Optimization
 flags can be passed via `--copt` to the `bazel run` command above.
 
 To test the created package, first install the resulting wheel file:
+
 ```bash
 pip install /tmp/tensorflow_compression/tensorflow_compression-*.whl
 ```
@@ -230,13 +258,15 @@ Then run the unit tests (Do not run the tests in the workspace directory where
 `WORKSPACE` of `tensorflow_compression` repo lives. In that case, the Python
 interpreter would attempt to import `tensorflow_compression` packages from the
 source tree, rather than from the installed package system directory):
+
 ```bash
 pushd /tmp
 python -m tensorflow_compression.all_tests
 popd
 ```
 
 When done, you can uninstall the pip package again:
+
 ```bash
 pip uninstall tensorflow-compression
 ```

build_pip_pkg.py

@@ -64,7 +64,7 @@ def main(srcdir):
       packages=setuptools.find_packages(),
       install_requires=[
           'scipy >= 1',
-          'tensorflow_probability >= 0.9',
+          'tensorflow_probability >= 0.12.1',
       ],
       script_args=['sdist', 'bdist_wheel'],
       # Add in any packaged data.
@@ -78,7 +78,7 @@ def main(srcdir):
           'Intended Audience :: Education',
           'Intended Audience :: Science/Research',
           'License :: OSI Approved :: Apache Software License',
-          'Programming Language :: Python :: 3.5',
+          'Programming Language :: Python :: 3',
           'Programming Language :: Python :: 3.6',
           'Programming Language :: Python :: 3.7',
           'Programming Language :: Python :: 3.8',

requirements.txt

@@ -1,2 +1,2 @@
 scipy >= 1
-tensorflow_probability >= 0.9
+tensorflow_probability >= 0.12.1

tensorflow_compression/__init__.py

@@ -17,10 +17,10 @@
 try:
   import tensorflow as _tensorflow
   _tf_version = [int(v) for v in _tensorflow.version.VERSION.split(".")]
-  assert _tf_version[0] == 2 and _tf_version[1] == 3
+  assert _tf_version[0] == 2 and _tf_version[1] == 4
 except (ImportError, AssertionError):
   raise RuntimeError(
-      "For tensorflow_compression, please install TensorFlow 2.3.")
+      "For tensorflow_compression, please install TensorFlow 2.4.")
 
 # pylint: disable=wildcard-import
 from tensorflow_compression.python.distributions.deep_factorized import *