如何在TensorFlow数据集批次中按指定列分组提取最大值并做增量运算

你只需要修改split_feature_label函数中的最大值计算逻辑即可，调整后可以实现按前三列、最后两列分组求全局最大值再叠加的需求。

完整修改后代码 #

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, -6],
         [7, 7, 8, -7, -7],
         [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 print_dataset(ds):
    for inputs, targets in ds:
        print("---Batch---")
        print("Feature:", inputs.numpy())
        print("Label:", targets.numpy())
        print("")
 
    
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 split_feature_label(x):
        # 计算前三列全局最大值（跨时间步、跨前三列所有元素）
        max_group1 = tf.reduce_max(x[:, :, :3], axis=[1,2], keepdims=True)
        # 计算最后两列全局最大值（跨时间步、跨最后两列所有元素）
        max_group2 = tf.reduce_max(x[:, :, 3:], axis=[1,2], keepdims=True)
        # 构造叠加张量：前三列统一加max_group1，后两列统一加max_group2
        add_tensor = tf.concat([
            tf.repeat(max_group1, 3, axis=-1),
            tf.repeat(max_group2, 2, axis=-1)
        ], axis=-1)
        # 拆分特征标签并叠加对应值
        return x[:, :input_sequence_length, :] + add_tensor, x[:, input_sequence_length:, label_slice] + add_tensor[:, input_sequence_length:, :]
         
    feature_ds = feature_ds.map(split_feature_label)
     
    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=1)
print_dataset(ds)

逻辑说明 #

1. 用tf.reduce_max指定axis=[1,2]同时对时间步维度、特征维度求最大值，得到分组后的全局最大值
2. 通过tf.repeat将最大值扩展到对应分组的所有列，再拼接成和原特征维度一致的叠加张量
3. 张量广播机制会自动将叠加张量匹配到每个时间步的对应位置，完成叠加操作
   运行后第一个批次的输出完全符合你给出的预期结果。

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