具有多个输出的 TensorFlow AutoGraph 多项式模型

Question

我有一个张量流模型，其输出对应于多个多项式的系数。请注意，我的模型实际上还有另一组输出（多输出），但我在下面仅通过返回除多项式系数之外的输入来模拟这一点。

我在模型训练过程中遇到了很多麻烦，与张量形状有关。我已经验证该模型能够预测样本输入，并且损失函数适用于样本输出。但是，在训练期间，它会立即抛出错误（见下文）

对于每个输入，模型接受一个固定嵌入大小的输入，并输出2 次数为2 的多项式的系数。例如，单个输入的输出可能如下所示：

[array([[[1, 2,  3],
        [ 4,  5,  6]]]),
[...]]

对应于多项式[1*x^2+2*x+3, 4*x^2+5*x+6]。请注意，我隐藏了第二个输出。

我注意到tf.math.polyval 需要一个系数列表，这使得它在使用 AutoGrad 时变得不稳定。所以，我用纯张量实现了我自己的霍纳算法版本。

import numpy as np
import tensorflow as tf
import logging
import tensorflow.keras as K

@tf.function
def tensor_polyval(coeffs, x):
    """
    Calculates polynomial scalars from tensor of polynomial coefficients
    Tensorflow tf.math.polyval requires a list coeff, which isn't compatible with autograd

    # Inputs:
      - coeffs (NxD Tensor): each row of coeffs corresponds to r[0]*x^(D-1)+r[1]*x^(D-2)...+r[D]
      - x: Scalar!

    # Output:
      - r[0]*x^(D-1)+r[1]*x^(D-2)...+r[D] for row in coeffs
    """
    p = coeffs[:, 0]
    for i in range(1,coeffs.shape[1]):
      tf.autograph.experimental.set_loop_options(
        shape_invariants=[(p, tf.TensorShape([None]))])
      c = coeffs[:, i]
      p = tf.add(c, tf.multiply(x, p))
    return p

@tf.function
def coeffs_to_poly(coeffs, n):
    # Converts a NxD array of coefficients to N evaluated polynomials at x=n
    return tensor_polyval(coeffs, tf.convert_to_tensor(n))

下面是我的模型、损失函数和训练例程的超级简化示例：

def model_init(embedDim=8, polyDim=2,terms=2):
  input = K.Input(shape=(embedDim,))
  x = K.layers.Reshape((embedDim,))(input)
  aCoeffs = K.layers.Dense((polyDim+1)*terms, activation='tanh')(x)
  aCoeffs = K.layers.Reshape((terms, polyDim+1))(aCoeffs)

  model = K.Model(inputs=input, outputs=[aCoeffs, input])
  return model

def get_random_batch(batch, embedDim, dtype='float64'):
  x = np.random.randn(batch, embedDim).astype(dtype)
  y = np.array([1. for i in range(batch)]).astype(dtype)
  return [x, 
          y]

@tf.function
def test_loss(y_true, y_pred, dtype=dataType):
  an = tf.vectorized_map(lambda y_p: coeffs_to_poly(y_p[0],
                                                    tf.constant(5,dtype=dataType)),
                         y_pred)
  
  return tf.reduce_mean(tf.reduce_mean(an,axis=-1))

embedDim=8
polyDim=2
terms=2
dataType = 'float64'
tf.keras.backend.set_floatx(dataType)

model = model_init(embedDim, polyDim, terms)

XTrain, yTrain  = get_random_batch(batch=128,
                                  embedDim=embedDim)

# Init Model
LR = 0.001
loss = test_loss
epochs = 5
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=LR), loss=loss)

hist = model.fit(XTrain,
        yTrain,
        batch_size=4,
        epochs=epochs,
        max_queue_size=10, workers=2, use_multiprocessing=True)

我得到的错误与tensor_polyval函数有关：

    <ipython-input-15-f96bd099fe08>:3 test_loss  *
        an = tf.vectorized_map(lambda y_p: coeffs_to_poly(y_p[0],
    <ipython-input-5-7205207d12fd>:23 coeffs_to_poly  *
        return tensor_polyval(coeffs, tf.convert_to_tensor(n))
    <ipython-input-5-7205207d12fd>:13 tensor_polyval  *
        p = coeffs[:, 0]
    ...
    ValueError: Index out of range using input dim 1; input has only 1 dims for '{{node strided_slice}} = StridedSlice[Index=DT_INT32, T=DT_DOUBLE, begin_mask=1, ellipsis_mask=0, end_mask=1, new_axis_mask=0, shrink_axis_mask=2](coeffs, strided_slice/stack, strided_slice/stack_1, strided_slice/stack_2)' with input shapes: [3], [2], [2], [2] and with computed input tensors: input[3] = <1 1>.

令人沮丧的是，我完全能够使用模型对样本输入进行预测并计算样本损失：

test_loss(yTrain[0:5],
          model.predict(XTrain[0:5]),
          dtype=dataType)

运行良好。

在test_loss 函数中，特别是我指的是第一个输出，通过y_p[0]。它尝试计算n=5 处的多项式值，然后输出所有内容的平均值（同样，这只是模拟代码）。据我了解，y_p[1] 将引用第二个输出（在这种情况下，是输入的副本）。我认为tf.vectorized_map 应该在模型批次的所有输出中运行，但它似乎在切割一个额外的维度？？

我注意到，如果我删除模型中的输出,input（使其成为单个输出）并将test_loss 中的y_p[0] 更改为y_p，代码确实会进行训练。我不知道为什么添加额外输出时它会损坏，因为我对 tf.vectorized_map 的理解意味着它分别作用于 list y_pred

的每个元素

Answer 1

如果我们需要单个损失函数来同时接收多个输出，也许我们可以将它们连接在一起形成一个输出。

在这种情况下：

1. 模型结构的变化，这里我们打包输出：

def model_init(embedDim=8, polyDim=2, terms=2):
    input = K.Input(shape=(embedDim, ))
    x = K.layers.Reshape((embedDim, ))(input)
    aCoeffs = K.layers.Dense((polyDim + 1) * terms, activation='tanh')(x)
    # pack the two outputs, add flatten layers if their shapes are not batch*K
    outputs = K.layers.Concatenate()([aCoeffs, input])
    
    model = K.Model(inputs=input, outputs=outputs)
    model.summary()
    return model

1. 损失函数的变化，这里我们解压输出：

# the loss function needs to know these
polyDim = 2
terms = 2

@tf.function
def test_loss(y_true, y_pred, dtype=dataType):
    """Loss function for flattened outputs."""
    
    # unpack multiple outputs
    offset = (polyDim + 1) * terms
    aCoeffs = tf.reshape(y_pred[:, :offset], [-1, terms, polyDim + 1])
    inputs = y_pred[:, offset:]
    
    print(aCoeffs, inputs)
    
    # do something with the two unpacked outputs, like below
    an = tf.vectorized_map(
        lambda y_p: coeffs_to_poly(y_p, tf.constant(5, dtype=dataType)),
        aCoeffs)

    return tf.reduce_mean(tf.reduce_mean(an, axis=-1))

请注意，损失函数依赖于对输出原始形状的了解来恢复它们。考虑子类化tf.keras.losses.Loss。

附：对于任何人只需要针对多重损失的不同损失：

1. 为两个输出定义损失函数。

@tf.function
def test_loss(y_true, y_pred, dtype=dataType):
    """Loss function for output 1
    (Only changed y_p[0] to y_p)"""
    an = tf.vectorized_map(
        lambda y_p: coeffs_to_poly(y_p, tf.constant(5, dtype=dataType)),
        y_pred)

    return tf.reduce_mean(tf.reduce_mean(an, axis=-1))


@tf.function
def dummy_loss(y_true, y_pred, dtype=dataType):
    """Loss function for output 2 i.e. the input, for debugging
    Better use 0 insead of 1.2345"""
    return tf.constant(1.2345, dataType)

1. 改成model.compile:

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=LR), loss=[test_loss, dummy_loss])