如何修改自定义tf.keras.Model使高维输入仍输出(1,3)张量？

我定义了一个继承自tf.keras.Model的Actor类：

import tensorflow as tf
from tensorflow.keras.layers import Dense

class Actor(tf.keras.Model):

    def __init__(self):
        super().__init__()
        self.linear1 = Dense(128, activation = 'relu')
        self.linear2 = Dense(256, activation = 'relu')
        self.linear3 = Dense(3, activation = 'softmax')

    # model override method        
    def call(self, state):
        x = tf.convert_to_tensor(state)
        x = self.linear1(x)
        x = self.linear2(x)
        x = self.linear3(x)
        return x

使用方式如下：

prob = self.actor(np.array([state]))

当输入维度为(5,)的state时，模型返回预期的(1,3)张量：

state: (5,) data: [0.50267935 0.50267582 0.50267935 0.50268406 0.5026817 ]
prob: (1, 3) data: tf.Tensor([[0.29540768 0.3525798  0.35201252]], shape=(1, 3), dtype=float32)

但输入维度为(5,3)的state时，模型返回(1,5,3)张量，不符合预期：

state: (5, 3) data: [[0.50789109 0.49648439 0.49651666]
 [0.5078905  0.49648391 0.49648928]
 [0.50788815 0.49648356 0.49643452]
 [0.50788677 0.4964834  0.49640713]
 [0.50788716 0.49648329 0.49635237]]
prob: (1, 5, 3) data: tf.Tensor(
[[[0.34579638 0.342928   0.3112757 ]
  [0.34579614 0.34292707 0.31127676]
  [0.34579575 0.34292522 0.31127906]
  [0.3457955  0.3429243  0.31128016]
  [0.34579512 0.34292242 0.3112824 ]]], shape=(1, 5, 3), dtype=float32)

使用版本为TensorFlow 2.9.1、Keras 2.9.0，需要修改模型使其对(5,3)输入仍返回(1,3)张量。

解决方案 #

问题根源在于Dense层默认作用在输入的最后一个维度，当输入形状为(1,5,3)时，Dense层会为每个5维度下的3维特征单独计算，最终保留5这个维度，导致输出形状为(1,5,3)。要解决这个问题，需要先将输入的多维度特征合并为一维，或者对序列维度做池化处理，以下是两种可行方案：

方案1：使用Flatten层展平输入 #

在模型中加入Flatten层，将(1,5,3)的输入展平为(1,15)，后续Dense层处理后即可输出(1,3)：

import tensorflow as tf
from tensorflow.keras.layers import Dense, Flatten

class Actor(tf.keras.Model):

    def __init__(self):
        super().__init__()
        self.flatten = Flatten()  # 新增展平层
        self.linear1 = Dense(128, activation='relu')
        self.linear2 = Dense(256, activation='relu')
        self.linear3 = Dense(3, activation='softmax')

    def call(self, state):
        x = tf.convert_to_tensor(state)
        x = self.flatten(x)  # 先展平输入
        x = self.linear1(x)
        x = self.linear2(x)
        x = self.linear3(x)
        return x

方案2：使用全局池化层（保留特征维度意义） #

如果需要保留每个特征维度的全局信息，可以使用GlobalAveragePooling1D层，对序列维度（这里是5）做平均池化，将(1,5,3)转为(1,3)：

import tensorflow as tf
from tensorflow.keras.layers import Dense, GlobalAveragePooling1D

class Actor(tf.keras.Model):

    def __init__(self):
        super().__init__()
        self.global_pool = GlobalAveragePooling1D()  # 全局平均池化层
        self.linear1 = Dense(128, activation='relu')
        self.linear2 = Dense(256, activation='relu')
        self.linear3 = Dense(3, activation='softmax')

    def call(self, state):
        x = tf.convert_to_tensor(state)
        x = self.global_pool(x)  # 对序列维度做池化
        x = self.linear1(x)
        x = self.linear2(x)
        x = self.linear3(x)
        return x

两种方案修改后，输入(5,3)的state（通过np.array([state])传入，形状为(1,5,3)）时，模型都会返回形状为(1,3)的张量，符合需求。

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