TensorFlow自定义数据增强预处理层实现及张量赋值报错解决

TensorFlow的张量（包括EagerTensor）是不可变对象，不支持numpy式的切片原地赋值操作，你代码里的img[:,pos:(pos+length),:] = 0就是触发报错的直接原因。
另外你的代码还有一个隐藏问题：用Python标准库random生成随机参数，在模型静态图运行、序列化保存加载的场景下，随机值只会在函数第一次被追踪时生成一次，后续不会动态随机，数据增强会失效。

所有操作改用TensorFlow原生API实现，通过张量拼接生成新的掩码后张量，不修改原张量，同时保证随机逻辑可被TF框架追踪：

import tensorflow as tf
from tensorflow.keras import layers

def random_mask_time(img):
    MAX_OCCURRENCE = 5
    MAX_MASK_LENGTH = 10
    time_axis_len = tf.shape(img)[0]
    # 用TF原生随机接口生成本次掩码块数量
    mask_count = tf.random.uniform(
        shape=[], minval=0, maxval=MAX_OCCURRENCE + 1, dtype=tf.int32
    )

    def _loop_step(i, current_img):
        # 生成单个掩码块的长度、起始位置
        block_len = tf.random.uniform(
            shape=[], minval=0, maxval=MAX_MASK_LENGTH + 1, dtype=tf.int32
        )
        start = tf.random.uniform(
            shape=[], minval=0, maxval=time_axis_len - MAX_MASK_LENGTH, dtype=tf.int32
        )
        # 拼接非掩码区、0值掩码块生成新张量
        zero_block = tf.zeros(
            [block_len, tf.shape(img)[1], tf.shape(img)[2]], dtype=img.dtype
        )
        updated = tf.concat(
            [
                current_img[:start, :, :],
                zero_block,
                current_img[start + block_len:, :, :]
            ],
            axis=0
        )
        return i + 1, updated

    _, masked_img = tf.while_loop(
        cond=lambda i, _: i < mask_count,
        body=_loop_step,
        loop_vars=[0, img]
    )
    return masked_img

# 频率维度掩码逻辑和时间维度一致，仅操作轴改为频率轴（axis=1）
def random_mask_freq(img):
    MAX_OCCURRENCE = 5
    MAX_MASK_LENGTH = 8 # 可根据声谱图频率维度尺寸自行调整
    freq_axis_len = tf.shape(img)[1]
    mask_count = tf.random.uniform(
        shape=[], minval=0, maxval=MAX_OCCURRENCE + 1, dtype=tf.int32
    )

    def _loop_step(i, current_img):
        block_len = tf.random.uniform(
            shape=[], minval=0, maxval=MAX_MASK_LENGTH + 1, dtype=tf.int32
        )
        start = tf.random.uniform(
            shape=[], minval=0, maxval=freq_axis_len - MAX_MASK_LENGTH, dtype=tf.int32
        )
        zero_block = tf.zeros(
            [tf.shape(img)[0], block_len, tf.shape(img)[2]], dtype=img.dtype
        )
        updated = tf.concat(
            [
                current_img[:, :start, :],
                zero_block,
                current_img[:, start + block_len:, :]
            ],
            axis=1
        )
        return i + 1, updated

    _, masked_img = tf.while_loop(
        cond=lambda i, _: i < mask_count,
        body=_loop_step,
        loop_vars=[0, img]
    )
    return masked_img

# 构建数据增强流水线
data_augmentation = tf.keras.Sequential([
    layers.Lambda(random_mask_time, name="time_mask"),
    layers.Lambda(random_mask_freq, name="freq_mask"),
    layers.RandomCrop(input_shape[1], input_shape[0])
])

• 如果你不需要自定义掩码逻辑，可以直接使用TensorFlow IO内置的声谱图增强APItfio.audio.time_mask、tfio.audio.freq_mask，完全对齐SpecAugment论文的标准实现，稳定性更高。
• 如果后续需要把模型导出为SavedModel格式部署，建议把掩码逻辑封装为继承tf.keras.layers.Layer的自定义层，比Lambda层的序列化兼容性更好。

内容的提问来源于stack exchange，提问作者V.Hunon