TensorFlow数据集逐批次提取最大值并叠加到原批次的实现方法

你只需要修改数据集的map处理逻辑，先计算每个样本最后一列的最大值，再通过TensorFlow广播机制将最大值叠加到样本的所有元素上即可，实现代码如下：

import tensorflow as tf
import numpy as np
 
simple_data_samples = np.array([
         [1, 1, 1, 7, -1],
         [2, -2, 2, -2, -2],
         [3, 3, 3, -3, -3],
         [-4, 4, 4, -4, -4],
         [5, 5, 5, -5, -5],
         [6, 6, 6, -4, -1],
         [7, 7, 8, -7, -70],
         [8, 8, 8, -8, -8],
         [9, 4, 9, -9, -9],
         [10, 10, 10, -10, -10],
         [11, 5, 11, -11, -11],
         [12, 12, 12, -12, -12],
])

def single (ds):
    for x in ds:
        print(x)
    
def timeseries_dataset_multistep_combined(features, label_slice, input_sequence_length, output_sequence_length, sequence_stride, batch_size):
    feature_ds = tf.keras.preprocessing.timeseries_dataset_from_array(features, None, sequence_length=input_sequence_length + output_sequence_length, sequence_stride=sequence_stride ,batch_size=batch_size, shuffle=False)
     
    def add_sample_last_col_max(x):
        # 计算每个样本最后一列的最大值，输出形状为 (batch_size, 1)
        sample_max = tf.reduce_max(x[:,:,-1], axis=1, keepdims=True)
        # 扩展维度适配输入张量形状，触发广播将最大值叠加到所有元素上
        return x + tf.expand_dims(sample_max, axis=-1)
    
    feature_ds = feature_ds.map(add_sample_last_col_max)
    return feature_ds

ds = timeseries_dataset_multistep_combined(simple_data_samples, slice(None, None, None), input_sequence_length=4, output_sequence_length=2, sequence_stride=2, batch_size=2)
single(ds)

运行后输出的第一个批次第一个样本和你给出的预期结果完全一致。

内容的提问来源于stack exchange，提问作者freak11