python - 使用 tf.keras 功能 API 时如何在 tensorflow/keras 中使用 for 循环？

Question

我希望使用 for 循环修改 CuDNNGRU 输出。但是，由于 tf.GradientTape 图形模式，我似乎无法这样做。如何在功能 API 中修改 CuDNNGRU？我知道通常我们可以使用 tf.keras.backend.* 函数对函数 API 执行一些矩阵操作，例如 K.backend.batch_dot 等。但是，我必须执行一些复杂的操作，例如三重循环或更多循环等，如果有人知道该怎么做，请帮忙！

.....source code
x = L.Lambda(lambda fm: tf.squeeze(fm, axis=1))(x)
gru_1 = CuDNNGRU(512, return_sequences=True, name='gru1')(x)
gru_1b = CuDNNGRU(512, return_sequences=True, go_backwards=True,name='gru1_b')(x)
for i in gru_1:
    .....apply some function to gru_1 outputs

顺便说一句，我目前尝试使用以下代码修改 GRU 输出。

def attention(inputs):
    transpose_input = tf.transpose(inputs,perm=[0,2,1])
    atten_w = K.backend.batch_dot(inputs,transpose_input)
    atten_w = tf.linalg.set_diag(atten_w,tf.zeros(tf.shape(atten_w)[0:-1],dtype=tf.float32))
    atten_w = tf.nn.softmax(atten_w,axis=1)
    atten_v = tf.py_function(calculate_atten,inp=[inputs,atten_w],Tout=[tf.float32])
    atten_v = tf.convert_to_tensor(atten_v)
    atten_v.set_shape(self.input_shapex)


def calculate_atten(data,atten_w):
    input_vector = data.numpy()
    atten_vectors = atten_w.numpy()
    all_batch = []
    for index,one_batch in enumerate(input_vector):
        tmp_w = atten_vectors[index]
        all_vector = []
        for j,vector in enumerate(one_batch):
            tmp = np.zeros(input_vector.shape[2])
            for w in tmp_w[j]:
                tmp += vector*w
            all_vector.append(tmp)
        all_batch.append(all_vector)
    return all_batch

但是，上面的代码，tf.py_function 返回 [time,features] 而不是 [batch,time,features]，如果可以做到，我可以使用 tf.py_function 来构建一个层。但是好像不行，求助！！！

更新

我已经能够通过嵌套的 tf.map_fn 实现操作。尽管它需要正确考虑将什么传递给 tf.map_fn（必须在多个输出中返回多个输入）。希望这可以帮助其他人

def attn_transformation(inputs):
    inputs_transpose = tf.transpose(inputs)
    atten_w = tf.tensordot(inputs,inputs_transpose,axes=1)
    def transform(data):
        multiply_data = data[0]*data[1][...,tf.newaxis]
        return [multiply_data,data[1]]
    data = tf.map_fn(lambda x:transform(x),elems=([inputs,atten_w]))
    data = tf.reduce_sum(data[0],axis=1)
    return data

gru_1 = CuDNNGRU(512, return_sequences=True, name='gru1')(x)
gru_1b = CuDNNGRU(512, return_sequences=True, go_backwards=True,name='gru1_b')(x)
atten_vf = L.Lambda(lambda x: tf.map_fn(attn_transformation,x))(gru_1)

标签： pythontensorflowkerastf.keras

Answer 1

对于任何你想将它应用到i你tensor可以使用的任意操作tf.map_fn()

因此，例如，我们可以执行以下操作：

inp = Input(shape=(2,3))
gru = CuDNNGRU(512, return_sequences=True)(inp)

def dummy_operation_to_be_applied(row):
  return row + 1

out = Lambda(lambda x: tf.map_fn(dummy_operation_to_be_applied, x))(gru)

更新：

请注意，我们也可以嵌套tf.map_fn()以映射较低维度的操作。

例如：

def nested_op(x):
  return tf.reduce_max(x) + x

def dummy_operation_to_be_applied(row):
  return tf.map_fn(nested_op, row)

out = Lambda(lambda x: tf.map_fn(dummy_operation_to_be_applied, x))(gru)