tf.one_hot内存占用过高崩溃，求张量场景下的替代或优化方案

你的内存崩溃问题根源在于显式生成大规模密集独热张量——当y_pred被reshape成一维后，再做tf.one_hot会生成形状为[batch_size * MAX_LEN, 64]的巨大张量，内存占用随批量和序列长度线性增长，很容易超出显存。

同时你的代码存在潜在错误：y_pred是模型输出的softmax概率分布（形状[batch_size, MAX_LEN, vocab_size]），直接将其cast为int64会把0-1之间的浮点数强制转为0，完全无法得到有效的预测类别索引，正确做法是先取argmax得到预测类别。

下面提供两种更高效的解决方案，完全避免显式独热编码：

方案一：直接基于概率分布计算MRR（推荐） #

MRR的核心是计算真实标签在预测概率排序中的倒数排名，完全不需要独热编码。重写MRR计算逻辑，直接用模型输出的概率张量计算：

修改后的代码 #

def mrr_metric(y_true, y_pred_probs):
    # y_true: 形状 [batch_size, MAX_LEN]，稀疏整数标签
    # y_pred_probs: 形状 [batch_size, MAX_LEN, vocab_size]，模型输出的softmax概率
    # 获取每个真实标签对应的预测概率
    true_probs = tf.gather(
        y_pred_probs, 
        tf.cast(y_true, tf.int32), 
        axis=-1, 
        batch_dims=2  # 匹配前两个维度（batch和序列长度）
    )
    # 计算每个位置中，概率大于等于真实标签概率的样本数，加1即为真实标签的排名（排名从1开始）
    ranks = tf.reduce_sum(tf.cast(y_pred_probs >= true_probs[..., tf.newaxis], tf.float32), axis=-1)
    # 计算倒数排名的平均值，即MRR
    mrr = tf.reduce_mean(1.0 / ranks)
    return mrr

def mrr_metric_fn(y_true, y_pred):
    # 直接传入模型输出的概率张量，无需reshape和独热
    return mrr_metric(y_true, y_pred)

# 模型创建和编译部分保持不变
def create_classifier_bert_model():    
    inputs = layers.Input((config.MAX_LEN,), dtype=tf.int64)
    sequence_output = pretrained_bert_model(inputs)
    outputs = layers.Dense(vectorize_layer.vocabulary_size(), activation="softmax")(sequence_output)
    classifer_model = keras.Model(inputs, outputs=[outputs], name="prediction")
    optimizer = keras.optimizers.Adam(learning_rate=config.LR)
    classifer_model.compile(
        optimizer=optimizer, 
        loss="sparse_categorical_crossentropy", 
        weighted_metrics=['sparse_categorical_accuracy', mrr_metric_fn]
    )
    return classifer_model

这个方案的优势：

• 完全避免生成大尺寸独热张量，内存占用仅和原始概率张量一致
• 基于概率计算排名比基于硬预测（argmax）更准确，能保留softmax的概率信息

方案二：用稀疏张量替代密集独热张量 #

如果必须保留基于类别索引的逻辑，可以用稀疏张量存储独热编码（仅存储非零元素），大幅降低内存占用：

修改后的代码 #

def mrr_metric_fn(y_true, y_pred):
    # 先获取预测类别索引（修正原来的错误）
    y_pred_idx = tf.argmax(y_pred, axis=-1)  # 形状 [batch_size, MAX_LEN]
    y_pred_revel = tf.reshape(y_pred_idx, [-1])  # 形状 [batch_size * MAX_LEN]
    
    # 创建稀疏独热张量，仅存储非零位置
    sparse_one_hot = tf.sparse.SparseTensor(
        indices=tf.stack([tf.range(tf.size(y_pred_revel)), y_pred_revel], axis=1),
        values=tf.ones(tf.size(y_pred_revel), dtype=tf.float32),
        dense_shape=[tf.size(y_pred_revel), 64]
    )
    # 如果原来的mrr_metric支持稀疏张量，直接传入；否则可转为密集张量（但仅在必要时）
    # sparse_one_hot_dense = tf.sparse.to_dense(sparse_one_hot)
    mrr = mrr_metric(y_true, sparse_one_hot)
    return mrr

注意：如果你的mrr_metric函数不支持稀疏张量，需要修改其逻辑适配稀疏输入，否则转成密集张量仍会有内存问题，因此方案一仍是最优选择。

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