在推理时更改保存的张量流模型输入形状

Change saved tensorflow model input shape at inference time

我找遍了所有地方,但一无所获。看起来很奇怪,没有人遇到过和我一样的问题......让我解释一下:

我训练了一个 Tensorflow 2 自定义模型。在训练期间,我使用了 set_shape((None, 320, 320, 14)) 以便 Tensorflow 知道形状(无论出于何种原因它都无法推断它...... -_-”)。 我还使用以下方法每 100 个周期保存了我的 自定义模型

model.save(os.path.join('models', 'pb', FLAGS.task_name + '-%i' % epoch))

所以对于第 100 个纪元,我将有一个文件夹 models/pb/my_name-100,其中包含

现在,在推理时,我只想加载模型(没有所有代码)。所以我创建了另一段代码,只加载模型并进行预测......基本模板如下:

class NeuralNetwork:
    def __init__(self, model):
        self.model = tf.keras.models.load_model(model)

    def predict(self, input_tensor):
        pred = self.model(input_tensor[None, ...])
        return pred[0]

其中 input_tensor 的尺寸为 (H, W, 14) 因此 input_tensor[None, ...] 的尺寸为: (None, 高, 宽, 14).

问题在于,因为我在训练期间将形状设置为 (None, 320, 320, 14)... 这个愚蠢的 Tensorflow 期望输入为 (None , 320, 320, 14) -_-”!!! 我的神经网络是一个全卷积神经网络,所以我真的不关心输入形状。我将它设置为 (320, 320, 14 ) 在训练期间出于记忆原因...

在预测期间,我希望能够对任何一种形状进行预测。

显然,我可以做一个预处理函数,从输入图像中提取大小为 (320, 320) 的补丁并将它们拼贴。例如,我的 input_tensor 的大小可能是 (30, 320, 320, 14)

然后在预测之后,我可以从图块中重建图像...但我不想那样做。

所以我的问题很简单: 我如何告诉 tensorflow 在推理时接受任何宽度和高度?天哪,这太麻烦了。我不敢相信没有简单的选择可以做到这一点

我回答我自己的问题。 不幸的是,我的回答 不会令 所有人满意。在 TF 中发生了太多令人费解的事情(更不用说当你寻求帮助时,大多数都是关于第一个 API...-_-")。

无论如何,这是“解决方案”

在我的神经网络中,我实现了一个 自定义层 来模仿 pytorch 函数 AdaptiveAvgPool2D。我的实现实际上在幕后使用 tf.nn_avg_pool 并且需要动态计算 kernel_size 以及 步幅 。这是我的代码,供参考:

class AdaptiveAvgPool2d(layers.Layer):
    def __init__(self, output_shape, data_format='channels_last'):
        super(AdaptiveAvgPool2d, self).__init__(autocast=False)
        assert data_format in {'channels_last', 'channels_first'}, \
            'data format parameter must be in {channels_last, channels_first}'

        if isinstance(output_shape, tuple):
            self.out_h = output_shape[0]
            self.out_w = output_shape[1]
        elif isinstance(output_shape, int):
            self.out_h = output_shape
            self.out_w = output_shape
        else:
            raise RuntimeError(f"""output_shape should be an Integer or a Tuple2""")

        self.data_format = data_format

    def call(self, inputs, mask=None):
        # input_shape = tf.shape(inputs)
        input_shape = inputs.get_shape().as_list()

        if self.data_format == 'channels_last':
            h_idx, w_idx = 1, 2
        else: # can use else instead of elif due to assert in __init__
            h_idx, w_idx = 2, 3

        stride_h = input_shape[h_idx] // self.out_h
        stride_w = input_shape[w_idx] // self.out_w

        k_size_h = stride_h + input_shape[h_idx] % self.out_h
        k_size_w = stride_w + input_shape[w_idx] % self.out_w

        pool = tf.nn.avg_pool(
            inputs,
            ksize=[k_size_h, k_size_w],
            strides=[stride_h, stride_w],
            padding='VALID',
            data_format='NHWC' if self.data_format == 'channels_last' else 'NCHW')

        return pool

问题是,我正在使用 inputs.get_shape().as_list() 恢复 int 值而不是 Tensor(..., type=int)。实际上,tf.nn.avg_pool 接受 ksizestrides 参数的 Int 列表...

换句话说,我不能使用 tf.shape(inputs) 因为它 returns 一个 Tensor(..., type=int) 并且除了评估它之外没有办法从 Tensor 中恢复一个 int .. .

我实现我的函数的方式工作得很好,问题是,Tensorflow 推断引擎盖下的大小并将所有张量的大小保存在 .pb 文件,当我保存它时。

事实上,您可以使用任何文本编辑器 (SublimeText) 轻松打开 .pb 文件,并自行查看预期的 TensorShape。在我的例子中是`TensorShape: [null, 320, 320, 14]

因此,使用 set_shape((None, None, None, 14)) 而不是 set_shape((None, 320, 320, 14))nothing 实际上并没有改变问题...

问题是平均池化层不接受动态内核size/strides...

然后我意识到实际上有一个 tensorflow 函数 tfa.layers.AdaptiveAveragePooling2D。所以,我可能会接受它,它会没事的,对吧?

不完全是,在幕后,这个 tensorflow 函数使用其他 tf.function,比如 tf.splittf.split 的问题在于,如果您要拆分的维度的大小为 X,并且您想要输出大小为 Y 的张量。如果 X % Y != 0,当 tf.split 将抛出一个错误...虽然 Pytorch 更强大并且处理案例是 X % Y != 0.

换句话说,这意味着,为了让我使用 tfa.layers.AdaptiveAveragePooling2D,我需要确保这个函数接收到的张量的大小可以被我传递给函数的标量整除.

例如,就我而言, 输入图像的大小为:(320, 320, whatever),tfa.layers.AdaptiveAveragePooling2D 接收到的输入张量为:(40, 40, whatever).

所以这意味着,我的张量的空间维度在训练过程中被除以 8。为了让它起作用,我应该选择一个可以除以 40 的大小。假设我选择 5.

这意味着在预测过程中,如果 tfa.layers.AdaptiveAveragePooling2D 接收的输入维度也能被 5 整除,我的神经网络将起作用。但是我们已经知道我的输入图像比张量接收到的 tfa.layers.AdaptiveAveragePooling2D 大 8 倍,所以这意味着,在预测时间,我可以使用任何图像大小只要:

  • H % (8 * 5) == 0 和 W % (8 * 5) == 0 其中 HW 分别是我的输入图像的高度和宽度。

为此,我们可以实现一个简单的函数: new_W = W + W % 40(本例中为 40...) new_H = H + H % 40(本例中为 40...)

此函数会稍微拉伸图像,但不会拉伸太多,应该没问题。

总结:

  • 我的 AdaptivePooling 使用静态形状,但我不能这样做,因为它在不接受动态形状的引擎盖下使用 tf.nn.avg_pool
  • tfa.layers.AdaptiveAveragePooling2D 是一种变通方法,但是因为它依赖于 tf.split,它对不精确的除法不稳健,所以它也不完美
  • 基本解决方案是使用 tfa.layers.AdaptiveAveragePooling2D 并在调用预测之前创建一个预处理函数,这样张量就可以在 tf.split 约束
  • 下正常工作
  • 最后,这也不是一个好的解决方案。因为,在训练期间,如果我收到一个大小为 (40, 40) 的张量并想要一个大小为 (5, 5)avg 输出,这意味着我基本上平均 (8, 8) features 检索一个特征。
  • 问题是,如果我在推理时间内对更大的图像执行此操作,我将收到更大的张量。比方说:(100, 200)。但是由于我的输出将始终是 (5, 5),这意味着这次我将平均 (20, 40) 个特征来检索 一个 个特征...

由于训练和推理之间的这种差异,如果我采用这种方式进行推理,则在更大的图像上进行推理可能会导致结果不一致 就我而言,方法是按照我在第一个 post...

中解释的那样对图像进行批处理

希望对大家有所帮助。