TensorFlow 2中含输入层信息的自定义损失函数训练问题

问题核心 #

你遇到的维度不匹配问题，本质是计算约束时输入张量与模型输出张量的批处理维度未对齐——批大小为32时，模型输出是(32,4)，但如果错误使用了整个训练集（320条）的输入张量，就会出现[32]与[320]的维度冲突。

可行解决方案 #

方案1：将约束融入自定义损失函数（推荐） #

把约束条件作为惩罚项加入损失函数，让模型在训练过程中自动趋近于满足等式。核心是确保所有运算在当前批维度上进行：

import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Input
import numpy as np

# 自定义损失：原回归损失 + 约束惩罚
def constrained_loss(y_true, y_pred):
    # 从y_true中提取当前批的输入x（需提前将输入嵌入标签）
    x = y_true[:, :2]
    x1, x2 = x[:, 0], x[:, 1]
    
    # 拆分模型输出的y1-y4
    y1, y2, y3, y4 = y_pred[:, 0], y_pred[:, 1], y_pred[:, 2], y_pred[:, 3]
    
    # 计算约束项，确保所有张量维度匹配（均为当前批大小）
    term1 = y1 / tf.sin(2 * x1)
    term2 = y2 / tf.cos(x1 + x2)
    term3 = y3 / tf.sin(x1 + x2)
    term4 = y4 / tf.cos(2 * x2)
    constraint = tf.square(term1 + term2 - term3 - term4)
    
    # 原回归任务损失（示例用MSE，根据你的任务调整）
    mse_loss = tf.keras.losses.MSE(y_true[:, 2:], y_pred)
    
    # 总损失：平衡原任务损失与约束惩罚，权重可自行调整
    return mse_loss + 10.0 * tf.reduce_mean(constraint)

# 构建基础网络
input_layer = Input(shape=(2,))
x = Dense(64, activation='relu')(input_layer)
x = Dense(32, activation='relu')(x)
output_layer = Dense(4)(x)
model = Model(inputs=input_layer, outputs=output_layer)

# 编译模型
model.compile(optimizer='adam', loss=constrained_loss)

# 生成模拟训练数据：将输入x嵌入标签前两列，后两列为真实y值
X = np.random.uniform(0.1, np.pi/2, (320, 2))  # 避开三角函数零点，防止除以0
y_true = np.random.uniform(-10, 10, (320, 4))
y_true[:, :2] = X

# 启动训练
model.fit(X, y_true, batch_size=32, epochs=10)

方案2：修改网络结构，天然满足约束 #

通过结构设计让输出自动符合约束，比如固定y4由y1-y3和输入计算得到，无需额外损失惩罚：

import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Input
import numpy as np

def build_constrained_model():
    input_layer = Input(shape=(2,))
    x = Dense(64, activation='relu')(input_layer)
    x = Dense(32, activation='relu')(x)
    
    # 先输出y1、y2、y3
    y_part = Dense(3)(x)
    y1, y2, y3 = y_part[:, 0:1], y_part[:, 1:2], y_part[:, 2:3]
    
    # 从输入中提取x1、x2，计算满足约束的y4
    x1, x2 = input_layer[:, 0:1], input_layer[:, 1:2]
    y4 = (y1 / tf.sin(2 * x1) + y2 / tf.cos(x1 + x2) - y3 / tf.sin(x1 + x2)) * tf.cos(2 * x2)
    
    # 拼接所有输出
    output_layer = tf.concat([y1, y2, y3, y4], axis=1)
    return Model(inputs=input_layer, outputs=output_layer)

# 编译训练，使用普通回归损失即可
model = build_constrained_model()
model.compile(optimizer='adam', loss='mse')

# 模拟数据
X = np.random.uniform(0.1, np.pi/2, (320, 2))
y_true = np.random.uniform(-10, 10, (320, 4))

model.fit(X, y_true, batch_size=32, epochs=10)

关键注意事项 #

• 避开三角函数零点：训练数据要避免2x1=kπ、x1+x2=π/2+kπ、2x2=π/2+kπ等情况，否则会触发除以0的错误。
• 维度一致性：所有张量运算必须保证批处理维度一致（即当前批的大小，如32），禁止在损失函数或网络层中引用整个训练集的张量。
• 惩罚权重调整：方案1中的约束惩罚权重（如示例中的10.0）需要根据任务需求调整，确保原任务损失与约束之间的平衡。

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