如何使用TensorFlow为RNN和LSTM模型计算AUC并生成ROC曲线？

针对你正在使用的trainDNN函数，我会一步步说明如何添加AUC计算和ROC曲线生成的功能，分为模型输出调整、训练中跟踪AUC、生成ROC曲线三个部分：

1. 调整模型输出，获取预测概率 #

因为ROC曲线需要基于预测的概率值（而非硬分类结果），所以首先要确保你的模型输出正类的概率：

在trainDNN函数里，当定义完RNN/LSTM的输出后，添加softmax层并提取正类概率：

def trainDNN(path, n_days, n_features, n_neurons, train_sequences, train_lengths, train_y, test_sequences, test_y, test_lengths, lstm=False, n_epochs=50, batch_size=256, learning_rate=0.0003, TRAIN_REC=8, TEST_REC=8):
    n_outputs = 2
    # ... 你的模型初始化、RNN/LSTM层定义代码 ...
    
    # 假设你已经得到了RNN的输出（比如叫做rnn_output）
    logits = fully_connected(rnn_output, num_outputs=n_outputs, activation_fn=None)
    # 计算softmax概率，取正类（索引1）的概率
    y_proba = tf.nn.softmax(logits)[:, 1]
    
    # ... 你的损失函数定义（比如交叉熵） ...
    xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=train_y, logits=logits)
    loss = tf.reduce_mean(xentropy)

2. 在训练中计算并跟踪AUC #

你可以使用TensorFlow的内置指标来计算AUC，推荐用tf.keras.metrics.AUC()，它支持流式计算，适合在训练循环中更新：

# 在函数内初始化AUC指标
train_auc = tf.keras.metrics.AUC(name='train_auc')
test_auc = tf.keras.metrics.AUC(name='test_auc')

# 计算训练集和测试集的AUC更新操作
train_auc.update_state(train_y, y_proba)
test_auc.update_state(test_y, y_proba)

# 定义优化器
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
training_op = optimizer.minimize(loss)

# 初始化变量
init = tf.global_variables_initializer()

然后在训练循环里，每次迭代后更新并打印AUC：

with tf.Session() as sess:
    init.run()
    for epoch in range(n_epochs):
        # 打乱训练数据（你已经有shuffle代码）
        train_sequences, train_lengths, train_y = shuffle(train_sequences, train_lengths, train_y)
        
        # 分批训练
        for iteration in range(len(train_sequences) // batch_size):
            start_idx = iteration * batch_size
            end_idx = start_idx + batch_size
            feed_dict = {
                # 你的输入占位符，比如 sequences: train_sequences[start_idx:end_idx],
                # lengths: train_lengths[start_idx:end_idx],
                # y: train_y[start_idx:end_idx]
            }
            sess.run(training_op, feed_dict=feed_dict)
        
        # 每个epoch结束后计算并打印AUC
        train_auc_result = train_auc.result().eval()
        test_auc_result = test_auc.result().eval()
        print(f"Epoch {epoch+1}/{n_epochs}, Train AUC: {train_auc_result:.4f}, Test AUC: {test_auc_result:.4f}")
        
        # 重置指标，准备下一轮epoch
        train_auc.reset_states()
        test_auc.reset_states()

3. 训练完成后生成ROC曲线 #

ROC曲线需要用测试集的真实标签和预测概率来生成，这里我们结合sklearn.metrics和matplotlib来实现：

首先在函数末尾，训练完成后获取测试集的预测概率：

# 在Session内，训练完成后获取测试集的预测概率
test_probs = y_proba.eval(feed_dict={
    # 测试集的输入占位符，比如 sequences: test_sequences,
    # lengths: test_lengths
})

然后导入必要的库并绘制ROC曲线：

from sklearn.metrics import roc_curve, auc
import matplotlib.pyplot as plt

# 计算ROC曲线的假阳性率、真阳性率和阈值
fpr, tpr, thresholds = roc_curve(test_y, test_probs)
roc_auc = auc(fpr, tpr)

# 绘制ROC曲线
plt.figure()
plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (area = {roc_auc:.4f})')
plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('Receiver Operating Characteristic')
plt.legend(loc="lower right")
plt.show()

注意事项 #

• 确保你的train_y和test_y是整数标签（0或1），如果是one-hot编码，需要先转换为整数标签（比如用tf.argmax(y, axis=1)）。
• 如果使用旧版TensorFlow（<2.0），tf.keras.metrics.AUC同样可用，或者你也可以用tf.metrics.auc，但需要手动处理更新和重置。

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