Changes relative imports in models/ to package imports.

Relative imports don't work out of Colab notebooks because notebooks are
interactive and aren't run as a module. To work around this, we now prefix model
imports with the model name. This requires the model files to live in a
directory either in the Python path or in the current directory. Note that there
is a small chance that a model name clashes with a global module name, but this
risk seems unavoidable under the circumstances.

PiperOrigin-RevId: 350441806
Change-Id: I4a13fa8337b11b132142f23f89a7824db14ec88f

Author: Johannes Ballé

models/hific/archs.py

@@ -27,7 +27,7 @@
 
 import tensorflow_compression as tfc
 
-from .helpers import ModelMode
+from hific.helpers import ModelMode
 
 SCALES_MIN = 0.11
 SCALES_MAX = 256

models/hific/configs.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 """Configurations for HiFiC."""
 
-from . import helpers
+from hific import helpers
 
 
 _CONFIGS = {

models/hific/evaluate.py

@@ -29,9 +29,9 @@
 import tensorflow.compat.v1 as tf
 import tensorflow_compression as tfc
 
-from . import configs
-from . import helpers
-from . import model
+from hific import configs
+from hific import helpers
+from hific import model
 
 
 # Show custom tf.logging calls.

models/hific/hific_test.py

@@ -20,10 +20,10 @@
 import numpy as np
 import tensorflow.compat.v1 as tf
 
-from . import configs
-from . import helpers
-from . import model
-from . import train
+from hific import configs
+from hific import helpers
+from hific import model
+from hific import train
 
 
 class FakeHiFiC(model.HiFiC):

models/hific/model.py

@@ -21,11 +21,11 @@
 import tensorflow.compat.v1 as tf
 import tensorflow_datasets as tfds
 
-from . import archs
-from . import helpers
-from .helpers import ModelMode
-from .helpers import ModelType
-from .helpers import TFDSArguments
+from hific import archs
+from hific import helpers
+from hific.helpers import ModelMode
+from hific.helpers import ModelType
+from hific.helpers import TFDSArguments
 
 
 # How many dataset preprocessing processes to use.

models/hific/train.py

@@ -19,9 +19,9 @@
 
 import tensorflow.compat.v1 as tf
 
-from . import configs
-from . import helpers
-from . import model
+from hific import configs
+from hific import helpers
+from hific import model
 
 # Show custom tf.logging calls.
 tf.logging.set_verbosity(tf.logging.INFO)

models/toy_sources/ntc.py

@@ -5,7 +5,7 @@
 import tensorflow.compat.v2 as tf
 import tensorflow_probability as tfp
 import tensorflow_compression as tfc
-from . import compression_model
+from toy_sources import compression_model
 
 tfpd = tfp.distributions
 

models/toy_sources/toy_sources.ipynb

@@ -25,6 +25,26 @@
         "This code requires tensorflow-compression v2.x.\n"
       ]
     },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "cellView": "form",
+        "id": "wdA5NUZ-fSxG"
+      },
+      "source": [
+        "#@title Dependencies for Colab\n",
+        "\n",
+        "# Run this cell to install the necessary dependencies when running the notebook\n",
+        "# directly in a Colaboratory hosted runtime from Github.\n",
+        "\n",
+        "%tensorflow_version 2.x\n",
+        "!pip install tensorflow-compression>=2.0\n",
+        "![[ -e tfc ]] || git clone https://github.com/tensorflow/compression tfc\n",
+        "%cd tfc/models\n"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -88,7 +108,7 @@
       "source": [
         "#@title Source definitions\n",
         "\n",
-        "from . import sawbridge\n",
+        "from toy_sources import sawbridge\n",
         "\n",
         "\n",
         "def _rotation_2d(degrees):\n",
@@ -138,8 +158,8 @@
         "\n",
         "import pywt\n",
         "import scipy\n",
-        "from . import ntc\n",
-        "from . import vecvq\n",
+        "from toy_sources import ntc\n",
+        "from toy_sources import vecvq\n",
         "\n",
         "\n",
         "def _get_activation(activation, dtype):\n",

models/toy_sources/vecvq.py

@@ -1,7 +1,7 @@
 """Variational entropy-constrained vector quantization."""
 
 import tensorflow.compat.v2 as tf
-from . import compression_model
+from toy_sources import compression_model
 
 
 class VECVQModel(compression_model.CompressionModel):