制作自定义 Keras 层时不能使用未知的输入 Dims（批量大小）

Question

我正在尝试在 Keras（后端 TensorFlow）中构建一个自定义层，该层在卷积层的过滤器上执行 KMeans 聚类。在构建该层的逻辑时，我迭代了批处理大小，但 Keras/TensorFlow 似乎不允许这种情况发生，因为批处理大小在运行时之前是一个未知的维度。

我尝试跟踪错误消息，它引导我找到两个文件：keras/engine/training.py 和 keras/engine/training_utils.py，据我所知，该错误源于一个问题

ndim

变量被定义为

NoneType

，因为编译模型时没有已知的批量大小。

在浏览各种 StackOverflow 和 GitHub 资源时，我没有看到任何关于如何处理 Keras/TensorFlow 拒绝未知批量大小参数的解决方案。

作为参考，这里是示例代码：

import numpy as np
import tensorflow as tf
import keras
from sklearn.cluster import KMeans

class KMeansLayer(keras.layers.Layer):
    def __init__(self, clusters=8, n_init=5, trainable=False, **kwargs):
        super(KMeansLayer, self).__init__(**kwargs)
        self.clusters = clusters
        self.n_init = n_init

    def build(self, input_shape):
        self.output_s = (input_shape[0],input_shape[1], input_shape[2],1)
        self.depth = input_shape[3]
        self.built=True

    def call(self, inputs):

        def KMeansFunc(input_tens,clusters=self.clusters,n_init=self.n_init):
            base_mat = np.zeros((input_tens.shape[0],input_tens.shape[1],input_tens.shape[2]))

            for frame in range(input_tens.shape[0]):
                init_mat = np.zeros((input_tens.shape[1]*input_tens.shape[2]))
                # print(init_mat.shape)
                reshape_mat = np.reshape(input_tens[frame],(input_tens.shape[1]*input_tens.shape[2],input_tens.shape[3]))
                # print(reshape_mat.shape)
                kmeans_init = KMeans(n_clusters=clusters, n_init=n_init)
                class_pred = kmeans_init.fit_predict(reshape_mat)

                for clust in range(self.clusters):
                    init_mat[class_pred==clust] = np.mean(reshape_mat[class_pred==clust],axis=1)
                    init_mat[class_pred==clust] = np.mean(reshape_mat[class_pred==clust],None)
                # print(base_mat.shape)

                base_mat[frame]=np.reshape(init_mat,(input_tens.shape[1],input_tens.shape[2]))

            return np.expand_dims(base_mat,axis=-1).astype('float32')

            
        output = tf.py_func(KMeansFunc,[inputs],tf.float32) 
        return output

    def compute_output_shape(self, input_shape):
        return self.output_s


input_1 = keras.Input(shape=(28,28,1), name='input_1', dtype='float32')

conv_1 = keras.layers.Conv2D(filters=20, kernel_size=3, strides=1, padding='same', data_format='channels_last', activation='elu', kernel_initializer='glorot_uniform')(input_1)
pool_1 = keras.layers.MaxPooling2D(pool_size=2, padding='same', data_format='channels_last')(conv_1)

up_conv_1 = keras.layers.SeparableConv2D(filters=20, kernel_size=2, strides=1, padding='same', data_format='channels_last', activation='elu', kernel_initializer='glorot_uniform')(pool_1)
up_1 = keras.layers.UpSampling2D(size=(2, 2), interpolation='bilinear')(up_conv_1)
conv_2 = keras.layers.Conv2D(filters=20, kernel_size=3, strides=1, padding='same', data_format='channels_last', activation='elu', kernel_initializer='glorot_uniform')(up_1)

conv_3 = keras.layers.Conv2D(filters=3, kernel_size=3, strides=1, padding='same', data_format='channels_last', activation='elu', kernel_initializer='glorot_uniform')(conv_2)

kmeans_out = KMeansLayer(clusters=8,n_init=5)(conv_3)


model = keras.Model(inputs=[input_1], outputs=kmeans_out)
keras.utils.plot_model(model, show_shapes=True)
model.compile(optimizer='adam', loss='mse', metrics=['mse'])

从上面的代码中可以看出，如果我有一个大小为

(batch_size,28,28,3)

的自定义图层的输入，我希望创建大小为

(batch_size,28,28,1)

的输出。

运行上述代码得到的错误是：

Traceback (most recent call last):
  File "example_error_file.py", line 64, in <module>
    model.compile(optimizer='adam', loss='mse', metrics=['mse'])
  File "~/fenv/lib/python3.6/site-packages/keras/engine/training.py", line 347, in compile
    sample_weight, mask)
  File "~/fenv/lib/python3.6/site-packages/keras/engine/training_utils.py", line 426, in weighted
    axis=list(range(weight_ndim, ndim)))
TypeError: 'NoneType' object cannot be interpreted as an integer

我有两个主要问题：

• 在定义自定义 Keras 层时我是否做错了什么？
• 有什么方法可以强制 Keras 在不知道这种情况下的批量大小的情况下运行（这似乎很合理）？

TensorFlow版本：1.7.0 喀拉斯版本：2.2.4 Python：3.6.6

Answer 1

主要问题似乎是，如果我们只是在完全张量流中查看这个事件，因为直到运行时才知道形状，当张量流尝试编译此图时，整个循环都会出现一些错误 None 。

考虑使用 tf.function 或 autograph，而不是使用 tf.pyfunc（如果您使用的是 TensorFlow 版本 1.14.0 或更高版本）。

或者，您可以使用 tf.scan，它将一个函数应用于张量中的每个元素，其中每个元素由从维度 0 处的原始张量解压出来的每个张量组成。

两者都应该有效！