如何解决3D医疗数据用sample_weight配合tf.data.Dataset训练报错问题

报错根本原因 #

1. 样本权重形状不符合要求：你当前生成的weights形状为(54,48,48,80,4)，与one-hot编码的标签同维度，但Keras的加权分类损失期望每个体素对应单个标量权重，不需要匹配one-hot的类别通道维度，损失模块尝试挤压最后一维时发现维度为4而非1，直接触发报错。
2. 样本权重计算逻辑错误：你直接对one-hot编码的y_train做拉平计算权重，相当于把one-hot里的0/1都当做独立标签计算权重，完全不符合类别平衡的计算要求。
3. sample_weight_mode参数设置错误："temporal"模式仅适配时序类任务，期望权重维度为(batch, timesteps)，无法适配3D分割的逐体素加权场景。

修复步骤 #

步骤1：修正样本权重计算逻辑 #

先将one-hot编码的标签转为类别索引，再计算平衡权重，最终权重形状与3D输入空间维度匹配，去掉最后一维的类别通道：

import numpy as np
from sklearn.utils.class_weight import compute_sample_weight

# 将one-hot标签转为类别索引，形状变为(54,48,48,80)
y_train_idx = np.argmax(y_train, axis=-1)
# 基于类别索引计算平衡权重
weights = compute_sample_weight(class_weight='balanced', y=y_train_idx.flatten())
# 重塑为与空间维度匹配的形状，无最后一维类别通道
weights = weights.reshape(y_train_idx.shape)

步骤2：修正模型编译参数 #

删除错误的sample_weight_mode="temporal"参数，3D分割逐体素加权不需要显式指定该参数，Keras会自动匹配形状：

model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

步骤3：调整验证集输入（可选） #

如果需要对验证集也做加权损失计算，需要用同样逻辑生成验证集的样本权重，传入validation_data即可，否则保持原有配置即可。

更简洁的替代方案：使用自定义加权损失 #

不需要单独处理样本权重，直接定义带类别权重的损失函数，代码更简洁，也不会出现形状匹配问题：

import tensorflow as tf
from sklearn.utils.class_weight import compute_class_weight

# 计算每个类别的平衡权重
y_train_idx = np.argmax(y_train, axis=-1)
class_weights = compute_class_weight(
    class_weight='balanced',
    classes=np.unique(y_train_idx),
    y=y_train_idx.flatten()
)

# 自定义加权交叉熵损失
def weighted_categorical_crossentropy(weights):
    def loss(y_true, y_pred):
        epsilon = 1e-7
        y_pred = tf.clip_by_value(y_pred, epsilon, 1. - epsilon)
        cross_entropy = -y_true * tf.math.log(y_pred)
        weighted_entropy = weights * cross_entropy
        return tf.reduce_sum(weighted_entropy, axis=-1)
    return loss

# 编译模型时直接使用自定义损失
model.compile(optimizer='adam',
              loss=weighted_categorical_crossentropy(class_weights),
              metrics=['accuracy'])

使用该方案时不需要再生成和传入sample_weights，直接正常传入训练数据即可。

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