4. MirroredStrategy基础API：创建、scope与run方法

好，咱们直接进入正题。MirroredStrategy 是 TensorFlow 里最常用的分布式策略之一，尤其适合单机多卡场景。不过别被名字骗了——它其实也能做多机多卡，只是默认情况下是单机。

我个人习惯把 MirroredStrategy 理解为「镜像策略」。什么意思呢？就是每张 GPU 上都放一份完整的模型副本，数据被切成多份分给各卡。梯度呢？各卡算完自己的，然后同步平均，再更新所有副本的参数。说白了，就是「数据并行」的经典实现。

4.1 创建 MirroredStrategy 实例

创建策略对象是第一步。代码极其简单：

import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()
print('Number of devices: {}'.format(strategy.num_replicas_in_sync))

就这么两行，TensorFlow 会自动检测当前机器有多少张 GPU。我遇到过一个问题：有次在服务器上跑，明明插了 4 张卡，结果只检测到 1 张。后来发现是 CUDA_VISIBLE_DEVICES 环境变量没设对。嗯，这里要注意——如果你只想用其中几张卡，可以手动指定：

strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1", "/gpu:2"])

或者通过环境变量：

# 在运行脚本前设置
# export CUDA_VISIBLE_DEVICES=0,1,2,3

你想想看，如果机器上有 8 张卡，但你不小心把 8 张全用了，别人就没法用了。所以我在生产环境里，一般会先检查一下可用设备，再决定策略用几张卡。

4.2 scope 的作用——把模型「装进去」

创建好策略对象后，下一步就是把模型和训练逻辑放到策略的 scope 里。scope 是什么？说白了就是一个上下文管理器，告诉 TensorFlow：「这段代码里的变量创建、计算图构建，都要按照分布式的方式来处理。」

看代码：

with strategy.scope():
    # 在这里创建模型、优化器、损失函数
    model = create_model()
    optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()

为什么一定要放在 scope 里？因为 MirroredStrategy 需要把变量创建在每张 GPU 上，并且保证它们同步。如果不放进去，变量就只会在默认设备上创建，那就失去了分布式的意义。

我曾经犯过一个低级错误：模型放在 scope 里了，但优化器忘了放。结果训练时梯度更新不一致，loss 死活降不下去。排查了半天才发现是优化器变量没被镜像。所以我的建议是——所有可训练的东西，都放在 scope 里创建。

4.3 strategy.run()——真正的执行入口

模型和优化器都准备好了，接下来就是训练。在 MirroredStrategy 里，你不能直接调用 model.fit() 或者自己写 for 循环来训练。你需要用 strategy.run() 来分发计算。

先定义一个训练步骤函数：

def train_step(inputs, labels):
    with tf.GradientTape() as tape:
        predictions = model(inputs, training=True)
        loss = loss_fn(labels, predictions)
    
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

然后通过 strategy.run() 来执行：

per_replica_losses = strategy.run(train_step, args=(batch_inputs, batch_labels))

这里有个关键点：strategy.run() 返回的是一个 PerReplica 对象，里面包含了每张卡上的 loss 值。你不能直接拿这个值来打印或做 loss 监控，需要先聚合一下：

loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)

或者用 tf.reduce_mean 取平均。我个人习惯用 ReduceOp.SUM 然后除以卡数，这样更直观。

4.4 完整示例：一个简单的训练循环

把上面三部分串起来，就是一个完整的分布式训练流程：

import tensorflow as tf

# 1. 创建策略
strategy = tf.distribute.MirroredStrategy()
print('Using {} devices'.format(strategy.num_replicas_in_sync))

# 2. 在 scope 里构建模型
with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10)
    ])
    optimizer = tf.keras.optimizers.Adam()
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 3. 准备数据（这里用 tf.data 做分布式数据加载）
dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
dataset = dataset.batch(64)
dist_dataset = strategy.experimental_distribute_dataset(dataset)

# 4. 定义训练步骤
def train_step(inputs, labels):
    with tf.GradientTape() as tape:
        logits = model(inputs, training=True)
        loss = loss_fn(labels, logits)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

# 5. 训练循环
for epoch in range(10):
    total_loss = 0.0
    num_batches = 0
    
    for batch_inputs, batch_labels in dist_dataset:
        per_replica_loss = strategy.run(train_step, args=(batch_inputs, batch_labels))
        batch_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None)
        total_loss += batch_loss
        num_batches += 1
    
    print('Epoch {}, Loss: {}'.format(epoch, total_loss / num_batches))

这段代码虽然简单，但已经包含了 MirroredStrategy 的核心用法。你可能会问：为什么不用 model.fit()？其实 model.fit() 在 scope 里也能自动分布式训练，但如果你想精细控制每一步，或者做自定义训练逻辑，那 strategy.run() 就是你的不二之选。

4.5 避坑指南

最后分享几个我踩过的坑：

• 数据加载顺序：用 experimental_distribute_dataset 时，数据会被自动分片。如果你自己手动分片，可能会造成数据重复或遗漏。我建议直接用 TensorFlow 提供的分布式数据集 API。
• Batch size 的设定：你传给 dataset.batch() 的 size 是全局的。比如你设 64，有 4 张卡，那每张卡实际拿到的 batch size 是 16。这个要算清楚，不然模型收敛会出问题。
• 变量初始化：所有变量必须在 scope 里创建。我曾经在 scope 外创建了一个计数器变量，结果训练时它只在主卡上更新，其他卡上的副本根本没变。嗯，这种 bug 特别难查。

好了，MirroredStrategy 的基础 API 就这些。说白了就是三步：创建策略、用 scope 包裹模型、用 run 执行训练。下一章我们会深入 experimental_distribute_dataset 的细节，看看数据到底是怎么分片的。