Refactor VQ-VAE example to work with recent versions of tensorflow-pr…

…obability (keras-team#921) * refactor tfp usage to 0.14.0+ versions * refactor to work with tfp 0.14.0+ * refactor to work with tfp 0.14.0+ * pixel_cnn call returns a tensor, not a layer * pixel_cnn call returns a tensor, not a layer * pixel_cnn call returns a tensor, not a layer * Bump keras-cv package version to current in order to fix Docker build error * rename `layers` to `outputs` as returned by categorical layer Signed-off-by: reinvantveer <rein@vantveer.me>
JZakraoui · Jun 17, 2022 · ab67ab5 · ab67ab5
1 parent 5793bf6
commit ab67ab5
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 16 deletions.
diff --git a/examples/generative/ipynb/vq_vae.ipynb b/examples/generative/ipynb/vq_vae.ipynb
@@ -53,7 +53,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip install -q tensorflow-probability==0.13.0"
+    "!pip install -q tensorflow-probability"
    ]
   },
   {
@@ -739,10 +739,10 @@
    "source": [
     "# Create a mini sampler model.\n",
     "inputs = layers.Input(shape=pixel_cnn.input_shape[1:])\n",
-    "x = pixel_cnn(inputs, training=False)\n",
-    "dist = tfp.distributions.Categorical(logits=x)\n",
-    "sampled = dist.sample()\n",
-    "sampler = keras.Model(inputs, sampled)"
+    "outputs = pixel_cnn(inputs, training=False)\n",
+    "categorical_layer = tfp.layers.DistributionLambda(tfp.distributions.Categorical)\n",
+    "outputs = categorical_layer(outputs)\n",
+    "sampler = keras.Model(inputs, outputs)"
    ]
   },
   {

diff --git a/examples/generative/md/vq_vae.md b/examples/generative/md/vq_vae.md
@@ -36,7 +36,7 @@ TensorFlow Probability, which can be installed using the command below.
 
 
 ```python
-!pip install -q tensorflow-probability==0.13.0
+!pip install -q tensorflow-probability
 ```
 
 ---
@@ -795,10 +795,10 @@ them to our decoder to generate novel images.
 ```python
 # Create a mini sampler model.
 inputs = layers.Input(shape=pixel_cnn.input_shape[1:])
-x = pixel_cnn(inputs, training=False)
-dist = tfp.distributions.Categorical(logits=x)
-sampled = dist.sample()
-sampler = keras.Model(inputs, sampled)
+outputs = pixel_cnn(inputs, training=False)
+categorical_layer = tfp.layers.DistributionLambda(tfp.distributions.Categorical)
+outputs = categorical_layer(outputs)
+sampler = keras.Model(inputs, outputs)
 ```
 
 We now construct a prior to generate images. Here, we will generate 10 images.

diff --git a/examples/generative/vq_vae.py b/examples/generative/vq_vae.py
@@ -32,7 +32,7 @@
 """
 
 """shell
-pip install -q tensorflow-probability==0.13.0
+pip install -q tensorflow-probability
 """
 
 """
@@ -493,10 +493,10 @@ def call(self, inputs):
 
 # Create a mini sampler model.
 inputs = layers.Input(shape=pixel_cnn.input_shape[1:])
-x = pixel_cnn(inputs, training=False)
-dist = tfp.distributions.Categorical(logits=x)
-sampled = dist.sample()
-sampler = keras.Model(inputs, sampled)
+outputs = pixel_cnn(inputs, training=False)
+categorical_layer = tfp.layers.DistributionLambda(tfp.distributions.Categorical)
+outputs = categorical_layer(outputs)
+sampler = keras.Model(inputs, outputs)
 
 """
 We now construct a prior to generate images. Here, we will generate 10 images.

diff --git a/scripts/autogen.py b/scripts/autogen.py
@@ -41,7 +41,7 @@
 PROJECT_URL = {
     "keras": "https://github.com/keras-team/keras/tree/v2.9.0/",
     "keras_tuner": "https://github.com/keras-team/keras-tuner/tree/1.1.2/",
-    "keras_cv": "https://github.com/keras-team/keras-cv/tree/v0.2.6/",
+    "keras_cv": "https://github.com/keras-team/keras-cv/tree/v0.2.7/",
     "keras_nlp": "https://github.com/keras-team/keras-nlp/tree/v0.2.0/",
 }