如何交替运行神经网络，对一个变量运行一半的迭代，然后转到下一个目标类别的训练？

假设我们有一个三个目标类别的分类问题，并且我们想要交替地对这三个类别进行训练。具体地，我们想要对一个类别运行一半的迭代（称为“一半阶段”），然后转移到下一个目标类别并重复同样的过程。以下是实现此过程的示例代码：

import tensorflow as tf

# 准备数据集（这里假设数据集已经被加载并划分好了）
train_dataset = ...
val_dataset = ...

# 定义模型
model = tf.keras.models.Sequential([
   ... # 添加神经网络层
])

# 定义训练参数
initial_learning_rate = 0.1
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate, decay_steps=100000, decay_rate=0.96, staircase=True
)
optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()

# 定义训练函数
@tf.function
def train_one_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x)
        loss_value = loss_fn(y, logits)
    grads = tape.gradient(loss_value, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))
    return loss_value

# 开始训练
num_epochs = 10
batch_size = 32
num_half_iterations = 2

train_iter = iter(train_dataset)
val_iter = iter(val_dataset)

for epoch in range(num_epochs):
    # 交替训练每个目标类别
    for target_class in range(3):
        # 跑一半阶段
        for _ in range(num_half_iterations):
            x, y = next(train_iter)
            loss_value = train_one_step(x, y