导出含IndependentNormal层的训练签名时遇AttributeError错误

环境信息 #

tensorflow                2.15.
tensorflow-estimator      2.15.0,
tensorflow-io-gcs-filesystem 0.36.0,
tensorflow-probability    0.23.0
keras                     2.15.0
keras-preprocessing       1.1.2

问题现象 #

基于tfp.layers.IndependentNormal构建的Keras模型训练过程正常，但导出训练签名供C++ API调用时，触发以下错误：

AttributeError: in user code:

File "/home/aberberich/Shared/Andi/tf2Api/reworked/min_export_failure.py", line 57, in trainOp  *
    loss = nll(predictions, targets)
File "/home/aberberich/Shared/Andi/tf2Api/reworked/min_export_failure.py", line 37, in nll  *
    return -estimated_distribution.log_prob(targets)

AttributeError: 'SymbolicTensor' object has no attribute 'log_prob'

问题代码 #

import numpy as np
from tensorflow.keras.layers import Input, Dense
import tensorflow_probability as tfp
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
import tensorflow as tf

inputDim = 10
targetDim = 1


#build train data
samples = 1000

input_list = []
for ii in range(samples):
    input_list.append(np.arange(inputDim))

input_arr = np.array(input_list)
target = arr = np.random.normal(5.0, 0.5, (samples,1))

#define model
input = Input(shape=(inputDim))
distribution_params = Dense(2)(input)
outputs = tfp.layers.IndependentNormal(targetDim)(distribution_params)

#define loss
def nll(targets, estimated_distribution):
    return -estimated_distribution.log_prob(targets)

#compile and fit model
optimizer = Adam()
model = Model(inputs= [input] , outputs=[outputs])
model.compile(optimizer=optimizer, loss=nll)#, metrics = lossFunction)
model.summary()
model.fit(input_arr,target, shuffle=True, epochs=500)#, verbose = 2)

# test prediction
prediction = model(np.expand_dims(np.arange(inputDim), axis = 0))
print("prediction mean : ", prediction.mean())
print("stdDev = ", prediction.stddev())

#export training signature
@tf.function
def trainOp(inputs, targets):
    ### has to return loss ###
    with tf.GradientTape() as tape:
        predictions  = model(inputs)
        loss = nll(predictions, targets)
    gradients = tape.gradient(loss, model.trainable_variables)    
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

signatures = {}
signatures["trainOp"] = trainOp.get_concrete_function(inputs = tf.TensorSpec((None, inputDim), tf.float32), 
                                                      targets = tf.TensorSpec((None, targetDim), tf.float32))

model.save('./testExport/', save_traces = False, signatures = signatures)

错误原因 #

1. 参数顺序错误：trainOp中调用nll时，将predictions和targets的顺序写反。nll函数定义要求第一个参数是目标值、第二个是分布对象，但调用时传参顺序颠倒，导致将目标张量传入了需要分布对象的参数位置，触发log_prob属性不存在的错误。
2. 图模式输出类型不匹配：TFP的IndependentNormal层默认在图模式（tf.function包裹的代码）下返回张量而非分布对象，即使参数顺序正确，也无法调用分布特有的log_prob方法。

修复方案 #

1. 修正nll调用参数顺序 #

在trainOp中，将loss = nll(predictions, targets)改为：

loss = nll(targets, predictions)

2. 强制TFP层返回分布对象 #

修改模型定义中的IndependentNormal层，添加convert_to_tensor=False参数，确保无论Eager还是图模式都返回分布对象：

outputs = tfp.layers.IndependentNormal(targetDim, convert_to_tensor=False)(distribution_params)

3. 优化损失计算（可选） #

为了保证梯度计算的稳定性，建议对损失取均值：

loss = tf.reduce_mean(nll(targets, predictions))

修复后的完整代码 #

import numpy as np
from tensorflow.keras.layers import Input, Dense
import tensorflow_probability as tfp
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
import tensorflow as tf

inputDim = 10
targetDim = 1


#build train data
samples = 1000

input_list = []
for ii in range(samples):
    input_list.append(np.arange(inputDim))

input_arr = np.array(input_list)
target = np.random.normal(5.0, 0.5, (samples,1))

#define model
input = Input(shape=(inputDim))
distribution_params = Dense(2)(input)
# 强制返回分布对象而非张量
outputs = tfp.layers.IndependentNormal(targetDim, convert_to_tensor=False)(distribution_params)

#define loss
def nll(targets, estimated_distribution):
    return -estimated_distribution.log_prob(targets)

#compile and fit model
optimizer = Adam()
model = Model(inputs= [input] , outputs=[outputs])
model.compile(optimizer=optimizer, loss=nll)
model.summary()
model.fit(input_arr,target, shuffle=True, epochs=500)

# test prediction
prediction = model(np.expand_dims(np.arange(inputDim), axis = 0))
print("prediction mean : ", prediction.mean())
print("stdDev = ", prediction.stddev())

#export training signature
@tf.function
def trainOp(inputs, targets):
    with tf.GradientTape() as tape:
        predictions  = model(inputs)
        # 修正参数顺序，并计算损失均值
        loss = tf.reduce_mean(nll(targets, predictions))
    gradients = tape.gradient(loss, model.trainable_variables)    
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

signatures = {}
signatures["trainOp"] = trainOp.get_concrete_function(inputs = tf.TensorSpec((None, inputDim), tf.float32), 
                                                      targets = tf.TensorSpec((None, targetDim), tf.float32))

model.save('./testExport/', save_traces = False, signatures = signatures)

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