TensorFlow 2.17中Keras模型complex64无效报错及权重共享问题

一、TensorFlow 2.17 中 complex64 报错的原因与修复 #

• 原因：TensorFlow 2.17 对 Keras Input 层的复杂 dtype 处理逻辑更严格，早期版本（如2.13）存在宽松兼容性处理，允许直接在输入层指定tf.complex64，但新版本中默认限制了复杂 dtype 作为输入类型，因此触发Invalid dtype: complex64错误。
• 修复方案（两种可选）：
  1. 转换输入为实张量，在模型内重构复向量：
     将复向量拆分为实部、虚部拼接的实张量作为输入，模型内部再组合为复向量处理，适配新版本的 dtype 要求。
     代码修改示例：

     # 数据预处理：拆分复向量为实部+虚部拼接的实张量
     X_real = np.real(X)
     X_imag = np.imag(X)
     X_combined = np.concatenate([X_real, X_imag], axis=1)
     train_dataset = tf.data.Dataset.from_tensor_slices((X_combined, y))
     train_dataset = train_dataset.shuffle(len(X)).batch(batch_size)
     
     # 模型定义调整
     def ftlm_model(hdim):
         # 输入层改为float32类型，维度为原复向量的2倍
         inputs = layers.Input(shape=(hdim*2,), dtype=tf.float32)
         # 拆分实部和虚部并重构复张量
         in_real = inputs[:, :hdim]
         in_imag = inputs[:, hdim:]
         inputs_complex = tf.complex(in_real, in_imag)
     
         # 后续处理逻辑保持不变
         shared_dense = layers.Dense(hdim, activation='relu')
         dense_real = shared_dense(tf.math.real(inputs_complex))
         dense_imag = shared_dense(tf.math.imag(inputs_complex))
         
         out_man = output_manip()
         output = out_man(dense_real, dense_imag)
         model = models.Model(inputs=inputs, outputs=output)
         return model
     

  2. 保留复向量输入，调整输入层 dtype 声明方式：
     若需直接使用复向量输入，可通过tf.cast在模型内部强制转换类型，避免输入层直接声明complex64：

     def ftlm_model(hdim):
         # 输入层使用默认float32，后续转为complex64
         inputs = layers.Input(shape=(hdim,))
         inputs_complex = tf.cast(inputs, tf.complex64)
         
         in_real = tf.math.real(inputs_complex)
         in_imag = tf.math.imag(inputs_complex)
         
         # 后续逻辑保持不变
         shared_dense = layers.Dense(hdim, activation='relu')
         dense_real = shared_dense(in_real)
         dense_imag = shared_dense(in_imag)
         
         out_man = output_manip()
         output = out_man(dense_real, dense_imag)
         model = models.Model(inputs=inputs, outputs=output)
         return model
     

二、模型外拆分数据时的权重共享实现 #

若必须在模型外拆分实部、虚部数据，可通过预定义共享层实例，将其同时应用于两个输入分支，实现权重复用，具体代码示例：

# 模型外拆分数据
X_real = np.real(X)
X_imag = np.imag(X)
# 构建包含两个输入的数据集
train_dataset = tf.data.Dataset.from_tensor_slices(((X_real, X_imag), y))
train_dataset = train_dataset.shuffle(len(X)).batch(batch_size)

# 预定义共享层实例（关键：同一个实例用于两个分支）
shared_dense = layers.Dense(hdim, activation='relu')

# 构建多输入模型
def shared_weight_model(hdim):
    # 实部输入层
    input_real = layers.Input(shape=(hdim,), dtype=tf.float32)
    # 虚部输入层
    input_imag = layers.Input(shape=(hdim,), dtype=tf.float32)
    
    # 共享层分别处理两个输入分支，复用同一组权重
    dense_real = shared_dense(input_real)
    dense_imag = shared_dense(input_imag)
    
    # 合并输出（保持原逻辑）
    output = tf.reduce_sum(tf.square(dense_real) + tf.square(dense_imag))
    model = models.Model(inputs=[input_real, input_imag], outputs=output)
    return model

model = shared_weight_model(hdim)
model.summary()

# 编译与训练逻辑不变
custom_loss = CustomMSELoss()
model.compile(optimizer='adam', loss=custom_loss)
history = model.fit(train_dataset, epochs=epochs)

核心原理：共享层是同一个实例，因此在两个分支中调用时，会使用同一组可训练权重，训练过程中权重更新会同步作用于两个分支。

内容的提问来源于stack exchange，提问作者D-west