K-Means of Tensorflow - Graph disconnected error

Question

I am trying to write a function that runs KMeans on a dataset and outputs the cluster centroids. My aim is to use this in a custom keras layer, so I am using TensorFlow's implementation of KMeans that takes a tensor as the input dataset.

My problem however is that I can't make it work even as a standalone function. The problem comes from the fact that KMeans accepts a generator function that provides mini-batches instead of a plain tensor, but when I am using closure to do that, I get a graph disconnected error:

import tensorflow as tf                                           # version: 2.4.1
from tensorflow.compat.v1.estimator.experimental import KMeans

@tf.function
def KMeansCentroids(inputs, num_clusters, steps, use_mini_batch=False):
    # `inputs` is a 2D tensor

    def input_fn():
        # Each one of the lines below results in the same "Graph Disconnected" error. Tuples don't really needed but just to be consistent with the documentation
        return (inputs, None)
        return (tf.data.Dataset.from_tensor_slices(inputs), None)
        return (tf.convert_to_tensor(inputs), None)
            
    kmeans = KMeans(
            num_clusters=num_clusters,
            use_mini_batch=use_mini_batch)
        
    kmeans.train(input_fn, steps=steps)     # This is where the error happens
    return kmeans.cluster_centers()

>>> x = tf.random.uniform((100, 2))
>>> c = KMeansCentroids(x, 5, 10)

The exact error is:

ValueError: Tensor("strided_slice:0", shape=(), dtype=int32) must be from the same graph as Tensor("Equal:0", shape=(), dtype=bool) (graphs are FuncGraph(name=KMeansCentroids, id=..) and ).

If I were to use a numpy dataset and convert to tensor inside the function, the code would work just fine.
Also, making input_fn() return directly tf.random.uniform((100, 2)) (ignoring the inputs argument), would again work. That's why I am guessing that tensorflow doesn't support closures since it needs to build the computation graph at the beginning.
But I don't see how to work around that. Could it be a version error due to KMeans being a compat.v1.experimental module?

Note that the documentation of KMeans states for the input_fn():

The function should construct and return one of the following:

A tf.data.Dataset object: Outputs of Dataset object must be a tuple (features, labels) with same constraints as below.

A tuple (features, labels): Where features is a tf.Tensor or a dictionary of string feature name to Tensor and labels is a Tensor or a dictionary of string label name to Tensor. Both features and labels are consumed by model_fn. They should satisfy the expectation of model_fn from inputs.

user_9215780 · Answer

The problem you're facing is more about invoking tensor outside the created graph. Basically, when you called the .train function, a new graph will be created and that is with the graph defined in that input_fn and the graph defined in the model_fn.

kmeans.train(input_fn, steps=steps)

And, after that all the tensors those coming outside these functions will be treated as outsiders and won't part of this new graph. That's why you're getting a graph disconnected error for trying to use outsider tensor. To resolve this, you need to create the necessary tensors within these graphs.

import tensorflow as tf                                        
from tensorflow.compat.v1.estimator.experimental import KMeans

@tf.function
def KMeansCentroids(num_clusters, steps, use_mini_batch=False):
    def input_fn(batch_size):
        pinputs = tf.random.uniform((100, 2))
        dataset = tf.data.Dataset.from_tensor_slices((pinputs))
        dataset = dataset.shuffle(1000).repeat()
        return dataset.batch(batch_size)

    kmeans = KMeans(
            num_clusters=num_clusters,
            use_mini_batch=use_mini_batch)
    
    kmeans.train(input_fn = lambda: input_fn(5), 
                 steps=steps)     
    
    return kmeans.cluster_centers()

c = KMeansCentroids(5, 10)

Here is some more info for reading, 1. FYI, I tested your code with a few versions of tf > 2, and I don't think it's related to version error or something.

Re-mentioning here for future readers. An alternative of using KMeans within Keras layers:

K-Means of Tensorflow - Graph disconnected error

1 Answers