TensorFlow SavedModel output output had no dimensions when loading a model with BigQuery ML

Question

I'm trying to use BigQuery ML to load a saved tensorflow model to make predictions. However when I run the query that read the saved model in GCS I got the following error:

TensorFlow SavedModel output output had no dimensions, which is not supported

I saved a K-Means model using tf.compat.v1.estimator.experimental.KMeans and this is my code:

def input_fn():
    return tf.data.Dataset.from_tensors(tf.convert_to_tensor(df.loc[:,col_features], dtype=tf.float32)).repeat(1)

## Function for convergence of kmeans
def centroids_distance(num_clusters: int, centroids_old: np.array, centroids_new: np.array) -> float:
    return np.max([np.linalg.norm(centroids_old[r,:] - centroids_new[r,:]) for r in range(num_clusters)])

feature_columns = [tf.feature_column.numeric_column(feature_name, dtype=tf.float32) for feature_name in col_features]

previous_centers = np.zeros((k, len(col_features)))
kmeans = tf.compat.v1.estimator.experimental.KMeans(num_clusters=k,
                                                    seed=seed,
                                                    initial_clusters='kmeans_plus_plus',
                                                    distance_metric='squared_euclidean',
                                                    use_mini_batch=False,
                                                    feature_columns=feature_columns)
for s in range(max_iterations):
    kmeans.train(input_fn)
    cluster_centers = kmeans.cluster_centers()
    centroids_diff = centroids_distance(k_opt, previous_centers, cluster_centers)
    if centroids_diff < tol:
        break
    previous_centers = cluster_centers

inputFn = tf.estimator.export.build_parsing_serving_input_receiver_fn(tf.feature_column.make_parse_example_spec(feature_columns))
model_path = kmeans.export_saved_model("gs://my_bucket/my_folder/model", inputFn)

I'm quite new to Tensorflow and I don't understand what the error is pointing to. I checked in export_save_model if there were other args to set the output dimension but none of the seemed to be what I was looking for. Thanks in advance.

Answer 1

Find the same problem here: Is there a way to use a kmeans, tensorflow saved model in bigquery?.

The main issue is that the output tensor shape of TF built-in KMeans estimator model has unknown rank in the saved model.

Two possible ways to solve this:

Try training the KMeans model on BQML directly. Reimplement the TF KMeans estimator model to reshape the output tensor into a specific tensor shape.