tensorflow/keras "An op outside of the function building code is being passed a 'Graph' Tensor"

user_161792 0
Viewed 224

I'm a newcomer to Tensorflow/Keras, and I've been following along with the book Hands-On Machine Learning with Scikit-Learn and Tensorflow. Chapter 12 deals with customizing Tensorflow, and following along with the associated notebook (here) I've hit the following custom model:

class ReconstructingRegressor(keras.models.Model):
    def __init__(self, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.hidden = [keras.layers.Dense(30, activation="selu",
                                          kernel_initializer="lecun_normal")
                       for _ in range(5)]
        self.out = keras.layers.Dense(output_dim)

    def build(self, batch_input_shape):
        n_inputs = batch_input_shape[-1]
        self.reconstruct = keras.layers.Dense(n_inputs)
        super().build(batch_input_shape)
        
    def call(self, inputs, training=None):
        Z = inputs
        for layer in self.hidden:
            Z = layer(Z)
        reconstruction = self.reconstruct(Z)

        recon_loss = tf.reduce_mean(tf.square(reconstruction - inputs))        
        self.add_loss(0.05 * recon_loss)

        return self.out(Z)

When I go to train using this model, I get the following error:

TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
  @tf.function
  def has_init_scope():
    my_constant = tf.constant(1.)
    with tf.init_scope():
      added = my_constant * 2
The graph tensor has name: mul:0

The problem is self.add_loss(0.05 * recon_loss); after commenting that out everything runs fine. Presumably recon_loss is the "Graph" tensor and self.add_loss() is the op outside of the function building code, but -- if that applies to add_loss() -- I don't know how I would add to the loss from within call().

Full disclosure: I'm using Tensorflow 2.3 when the book was written with 2.1 in mind, so I'm not really following instructions. That said, I'm really curious how to fix this and at my current level of knowledge I feel essentially powerless. It seems like it should work -- how else to add to the loss function? Any help would be appreciated here.

Full example:

import tensorflow as tf
from tensorflow import keras

from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

housing = fetch_california_housing()
X_train_full, X_test, y_train_full, y_test = train_test_split(
    housing.data, housing.target.reshape(-1, 1), random_state=42)
X_train, X_valid, y_train, y_valid = train_test_split(
    X_train_full, y_train_full, random_state=42)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_valid_scaled = scaler.transform(X_valid)
X_test_scaled = scaler.transform(X_test)

class ReconstructingRegressor(keras.models.Model):
    def __init__(self, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.hidden = [keras.layers.Dense(30, activation="selu",
                                          kernel_initializer="lecun_normal")
                       for _ in range(5)]
        self.out = keras.layers.Dense(output_dim)

    def build(self, batch_input_shape):
        n_inputs = batch_input_shape[-1]
        self.reconstruct = keras.layers.Dense(n_inputs)
        super().build(batch_input_shape)
        
    def call(self, inputs, training=None):
        Z = inputs
        for layer in self.hidden:
            Z = layer(Z)
        reconstruction = self.reconstruct(Z)
        recon_loss = tf.reduce_mean(tf.square(reconstruction - inputs))
        
        self.add_loss(0.05 * recon_loss)

        return self.out(Z)

model = ReconstructingRegressor(1, dynamic=True)
model.compile(loss="mse", optimizer="nadam")
history = model.fit(X_train_scaled, y_train, epochs=2)
keras tensorflow python
1 Answers

Although I think it's too late to answer the question... Let me show you what I've tried.

First, I removed the build function of the custom model.

    def build(self, batch_input_shape):
        n_inputs = batch_input_shape[-1]
        self.reconstruct = keras.layers.Dense(n_inputs)
        super().build(batch_input_shape)

And run_eagerly = True when compiling OR using a custom layer calculating the custom loss worked. For example, write a custom layer code:

class ReconLoss(keras.layers.Layer):
  def __init__(self, **kwargs):
    super().__init__(**kwargs)

  def call(self, inputs):
    x, reconstruction = inputs
    recon_loss = tf.reduce_mean(tf.square(reconstruction - x))

    self.add_loss(0.05 * recon_loss)

    return

Then assign an instance of it in __init__ of the custom model, insert self.ReconLoss([x, reconstruction]) into call method of the custom model.

Edited colab code : https://colab.research.google.com/drive/1Hwi6auz2meKvD0ogdDSywb2E4_J1F9S_?usp=sharing

I don't understand still why the error was raised, but this worked for me.


References :
user_100345230
Related