Import of pull request #58: HiFiC for Camera Ready (author: fab-jul).

- Remove "PRE-RELEASE" warning.
- Improve evaluate.py:
- Do real arithmetic coding to report bitrate. This includes running the entropy model updates.
- Support image folders as input (not just TFDS).
- Report PSNR.
- Rename LayerNorm to ChannelNorm.
- Fix Colab to make it better to use:
- Support different HiFiC models with drop down.
- Support blindly running whole file.
- Support skipping the upload cell.
- support re-running with same files but potentially different model.
- Fix tf.logging statements not getting printed.
- Fix train.py doing back to back training.

PiperOrigin-RevId: 342352674
Change-Id: If3a78a5c37cf827aa0ed71a990c178144084b766

Author: relational

models/hific/README.md

@@ -1,7 +1,5 @@
 # High-Fidelity Generative Image Compression
 
-## PRE-RELEASE
-
 <div align="center">
   <a href='https://hific.github.io'>
   <img src='https://hific.github.io/social/thumb.jpg' width="80%"/>
@@ -96,7 +94,6 @@ If you get slow training/stalling, try tweaking the `DATASET_NUM_PARALLEL` and
 The architecture is defined in `arch.py`, which is used to build the model from
 `model.py`. Our configurations are in `configs.py`.
 
-
 We release a _simplified_ trainer in `train.py` as a starting point for custom
 training. Note that it's using
 [coco2014](https://cocodataset.org) from

models/hific/archs.py

@@ -61,7 +61,7 @@
     "HyperInfo",
     "decoded latent_shape hyper_latent_shape "
     "nbpp side_nbpp total_nbpp qbpp side_qbpp total_qbpp "
-    "bitstring side_bitstring",
+    "bitstream_tensors",
 )
 
 
@@ -86,7 +86,7 @@ def __init__(self,
     model = [
         tf.keras.layers.Conv2D(
             filters=num_filters_base, kernel_size=7, padding="same"),
-        LayerNorm(),
+        ChannelNorm(),
         tf.keras.layers.ReLU()
     ]
 
@@ -95,7 +95,7 @@ def __init__(self,
           tf.keras.layers.Conv2D(
               filters=num_filters_base * 2 ** (i + 1),
               kernel_size=3, padding="same", strides=2),
-          LayerNorm(),
+          ChannelNorm(),
           tf.keras.layers.ReLU()])
 
     model.append(
@@ -127,11 +127,11 @@ def __init__(self,
       num_filters_base: base number of filters.
       num_residual_blocks: number of residual blocks.
     """
-    head = [LayerNorm(),
+    head = [ChannelNorm(),
             tf.keras.layers.Conv2D(
                 filters=num_filters_base * (2 ** num_up),
                 kernel_size=3, padding="same"),
-            LayerNorm()]
+            ChannelNorm()]
 
     residual_blocks = []
     for block_idx in range(num_residual_blocks):
@@ -151,7 +151,7 @@ def __init__(self,
               filters=filters,
               kernel_size=3, padding="same",
               strides=2),
-          LayerNorm(),
+          ChannelNorm(),
           tf.keras.layers.ReLU()]
 
     # Final conv layer.
@@ -201,19 +201,19 @@ def __init__(
 
     block = [
         tf.keras.layers.Conv2D(**kwargs_conv2d),
-        LayerNorm(),
+        ChannelNorm(),
         tf.keras.layers.Activation(activation),
         tf.keras.layers.Conv2D(**kwargs_conv2d),
-        LayerNorm()]
+        ChannelNorm()]
 
     self.block = tf.keras.Sequential(name=name, layers=block)
 
   def call(self, inputs, **kwargs):
     return inputs + self.block(inputs, **kwargs)
 
 
-class LayerNorm(tf.keras.layers.Layer):
-  """Implement LayerNorm.
+class ChannelNorm(tf.keras.layers.Layer):
+  """Implement ChannelNorm.
 
   Based on this paper and keras' InstanceNorm layer:
     Ba, Jimmy Lei, Jamie Ryan Kiros, and Geoffrey E. Hinton.
@@ -238,7 +238,7 @@ def __init__(self,
       gamma_initializer: Initializer for gamma.
       **kwargs: Passed to keras.
     """
-    super(LayerNorm, self).__init__(**kwargs)
+    super(ChannelNorm, self).__init__(**kwargs)
 
     self.axis = -1
     self.epsilon = epsilon
@@ -478,6 +478,14 @@ def _make_synthesis(syn_name):
 
     self._side_entropy_model = FactorizedPriorLayer()
 
+  @property
+  def losses(self):
+    return self._side_entropy_model.losses
+
+  @property
+  def updates(self):
+    return self._side_entropy_model.updates
+
   @property
   def transform_layers(self):
     return [self._analysis, self._synthesis_scale, self._synthesis_mean]
@@ -529,7 +537,7 @@ def call(self, latents, image_shape, mode: ModelMode) -> HyperInfo:
 
     compressed = None
     if training:
-      latents_decoded = _quantize(latents, latent_means)
+      latents_decoded = _ste_quantize(latents, latent_means)
     elif validation:
       latents_decoded = entropy_info.quantized
     else:
@@ -546,16 +554,25 @@ def call(self, latents, image_shape, mode: ModelMode) -> HyperInfo:
         qbpp=entropy_info.qbpp,
         side_qbpp=side_info.total_qbpp,
         total_qbpp=entropy_info.qbpp + side_info.total_qbpp,
-        bitstring=compressed,
-        side_bitstring=side_info.bitstring)
+        # We put everything that's needed for real arithmetic coding into
+        # the bistream_tensors tuple.
+        bitstream_tensors=(compressed, side_info.bitstring,
+                           image_shape, latent_shape, side_info.latent_shape))
 
     tf.summary.scalar("bpp/total/noisy", info.total_nbpp)
     tf.summary.scalar("bpp/total/quantized", info.total_qbpp)
 
+    tf.summary.scalar("bpp/latent/noisy", entropy_info.nbpp)
+    tf.summary.scalar("bpp/latent/quantized", entropy_info.qbpp)
+
+    tf.summary.scalar("bpp/side/noisy", side_info.total_nbpp)
+    tf.summary.scalar("bpp/side/quantized", side_info.total_qbpp)
+
     return info
 
 
-def _quantize(inputs, mean):
+def _ste_quantize(inputs, mean):
+  """Calculates quantize(inputs - mean) + mean, sets straight-through grads."""
   half = tf.constant(.5, dtype=tf.float32)
   outputs = inputs
   outputs -= mean