Adds example training code for BMSHJ2018 model.

Also adds some minor tweaks to other example scripts.

PiperOrigin-RevId: 259822670
Change-Id: Iaf726302909190673eba2c174fa7c4989f7f5d3e

Author: Johannes Ball?

BUILD

@@ -39,6 +39,12 @@ py_binary(
     ],
 )
 
+py_binary(
+    name = "generate_docs",
+    srcs = ["tools/generate_docs.py"],
+    deps = [":tensorflow_compression"],
+)
+
 py_binary(
     name = "tfci",
     srcs = ["examples/tfci.py"],
@@ -52,7 +58,7 @@ py_binary(
 )
 
 py_binary(
-    name = "generate_docs",
-    srcs = ["tools/generate_docs.py"],
+    name = "bmshj2018",
+    srcs = ["examples/bmshj2018.py"],
     deps = [":tensorflow_compression"],
 )

examples/bls2017.py

@@ -22,6 +22,11 @@
 https://arxiv.org/abs/1611.01704
 
 With patches from Victor Xing <[email protected]>
+
+This is meant as 'educational' code - you can use this to get started with your
+own experiments. To reproduce the exact results from the paper, tuning of hyper-
+parameters may be necessary. To compress images with published models, see
+`tfci.py`.
 """
 
 from __future__ import absolute_import
@@ -176,11 +181,11 @@ def build(self, input_shape):
         tfc.SignalConv2D(
             self.num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
             padding="same_zeros", use_bias=True,
-            activation=tfc.GDN(name="gdn_0", inverse=True)),
+            activation=tfc.GDN(name="igdn_0", inverse=True)),
         tfc.SignalConv2D(
             self.num_filters, (5, 5), name="layer_1", corr=False, strides_up=2,
             padding="same_zeros", use_bias=True,
-            activation=tfc.GDN(name="gdn_1", inverse=True)),
+            activation=tfc.GDN(name="igdn_1", inverse=True)),
         tfc.SignalConv2D(
             3, (9, 9), name="layer_2", corr=False, strides_up=4,
             padding="same_zeros", use_bias=True,
@@ -290,7 +295,7 @@ def compress(args):
   y_hat, likelihoods = entropy_bottleneck(y, training=False)
   x_hat = synthesis_transform(y_hat)
 
-  num_pixels = tf.to_float(tf.reduce_prod(tf.shape(x)[:-1]))
+  num_pixels = tf.cast(tf.reduce_prod(tf.shape(x)[:-1]), dtype=tf.float32)
 
   # Total number of bits divided by number of pixels.
   eval_bpp = tf.reduce_sum(tf.log(likelihoods)) / (-np.log(2) * num_pixels)