TensorFlow/PyTorch：自定义训练循环与分布式数据并行的工程实践

二、框架差异的底层根源：静态图与动态图的哲学分歧

TensorFlow 和 PyTorch 的核心架构差异源于一个根本性的设计选择：计算图的构建时机。TensorFlow 1.x 采用静态图（Define-and-Run）——先定义完整计算图，再输入数据执行。PyTorch 采用动态图（Define-by-Run）——每次前向传播时动态构建计算图。这个选择影响了从调试体验到分布式训练的方方面面。

TensorFlow 2.x 引入了 Eager Mode 兼容动态图，但底层仍保留了静态图的优化能力（通过 tf.function 编译）。PyTorch 2.0 引入了 torch.compile 获得静态图优化能力，但默认仍是动态图。两个框架在"动态灵活性"和"静态优化"之间寻找各自的平衡点，而自定义训练循环和分布式数据并行是这种差异最明显的领域。

一、自定义训练循环：为什么 Model.fit() 不够用

框架提供的高级 API（model.fit() / tf.keras.Model.fit）覆盖了 80% 的标准训练场景，但剩余 20% 的需求往往是最关键的：混合精度训练中自定义 Loss Scaling、多模型交替训练（GAN）、梯度累积突破显存限制、自定义学习率调度与梯度裁剪的组合。这些场景要求开发者手动控制训练循环的每一步，而非依赖框架的"黑盒"封装。

自定义训练循环的代价是：必须手动处理所有细节——梯度清零、前向传播、损失计算、反向传播、参数更新、指标记录。任何一个步骤的遗漏都可能导致静默错误（如忘记 optimizer.zero_grad() 导致梯度累积）。

1.1 训练循环的统一抽象

flowchart TD
    A[数据加载] --> B[前向传播]
    B --> C[损失计算]
    C --> D[反向传播]
    D --> E[梯度处理<br/>裁剪/累积/缩放]
    E --> F[参数更新]
    F --> G[指标记录]
    G --> H{是否继续?}
    H -->|是| A
    H -->|否| I[训练结束]

    style D fill:#ffebee
    style E fill:#fff3e0
    style F fill:#e1f5fe

三、生产级代码实现：双框架对比

3.1 PyTorch 自定义训练循环

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.cuda.amp import GradScaler, autocast
from typing import Optional

class PyTorchTrainer:
    """PyTorch 自定义训练循环：支持混合精度与梯度累积"""

    def __init__(
        self,
        model: nn.Module,
        optimizer: torch.optim.Optimizer,
        scheduler=None,
        gradient_clip_norm: Optional[float] = 1.0,
        accumulation_steps: int = 1,
        use_amp: bool = True,
    ):
        self.model = model
        self.optimizer = optimizer
        self.scheduler = scheduler
        self.gradient_clip_norm = gradient_clip_norm
        self.accumulation_steps = accumulation_steps
        self.use_amp = use_amp

        # 混合精度训练的 GradScaler
        self.scaler = GradScaler(enabled=use_amp)

        # 梯度累积计数器
        self._accumulation_count = 0

    def train_epoch(
        self,
        train_loader: DataLoader,
        epoch: int,
        device: torch.device,
    ) -> dict:
        self.model.train()
        total_loss = 0.0
        total_samples = 0

        for step, (inputs, targets) in enumerate(train_loader):
            inputs = inputs.to(device, non_blocking=True)
            targets = targets.to(device, non_blocking=True)

            # 混合精度前向传播
            with autocast(enabled=self.use_amp):
                outputs = self.model(inputs)
                loss = nn.functional.cross_entropy(outputs, targets)
                # 梯度累积：损失除以累积步数
                loss = loss / self.accumulation_steps

            # 反向传播（缩放后的梯度）
            self.scaler.scale(loss).backward()

            self._accumulation_count += 1
            total_loss += loss.item() * self.accumulation_steps * inputs.size(0)
            total_samples += inputs.size(0)

            # 梯度累积达到步数后执行参数更新
            if self._accumulation_count % self.accumulation_steps == 0:
                # 梯度裁剪（在 unscale 之后）
                if self.gradient_clip_norm is not None:
                    self.scaler.unscale_(self.optimizer)
                    nn.utils.clip_grad_norm_(
                        self.model.parameters(),
                        self.gradient_clip_norm,
                    )

                # 参数更新
                self.scaler.step(self.optimizer)
                self.scaler.update()
                self.optimizer.zero_grad(set_to_none=True)

                # 学习率调度
                if self.scheduler is not None:
                    self.scheduler.step()

            # 定期打印训练状态
            if step % 100 == 0:
                avg_loss = total_loss / max(total_samples, 1)
                lr = self.optimizer.param_groups[0]["lr"]
                print(
                    f"Epoch {epoch} Step {step}: "
                    f"loss={avg_loss:.4f}, lr={lr:.2e}"
                )

        return {"train_loss": total_loss / max(total_samples, 1)}

3.2 TensorFlow 自定义训练循环

import tensorflow as tf
from typing import Optional

class TFTrainer:
    """TensorFlow 自定义训练循环：支持混合精度与梯度累积"""

    def __init__(
        self,
        model: tf.keras.Model,
        optimizer: tf.keras.optimizers.Optimizer,
        gradient_clip_norm: Optional[float] = 1.0,
        accumulation_steps: int = 1,
        use_amp: bool = True,
    ):
        self.model = model
        self.optimizer = optimizer
        self.gradient_clip_norm = gradient_clip_norm
        self.accumulation_steps = accumulation_steps
        self.use_amp = use_amp

        # 混合精度策略
        if use_amp:
            self.optimizer = tf.keras.mixed_precision.LossScaleOptimizer(
                optimizer
            )

        # 梯度累积缓冲区
        self._gradient_buffer = None
        self._accumulation_count = 0

    @tf.function
    def train_step(self, inputs, targets):
        """单步训练（编译为静态图以获得更好性能）"""
        with tf.GradientTape() as tape:
            outputs = self.model(inputs, training=True)
            loss = tf.keras.losses.sparse_categorical_crossentropy(
                targets, outputs, from_logits=True
            )
            loss = tf.reduce_mean(loss) / self.accumulation_steps

            # 混合精度缩放
            if self.use_amp:
                loss = self.optimizer.get_scaled_loss(loss)

        gradients = tape.gradient(loss, self.model.trainable_variables)

        # 混合精度反缩放
        if self.use_amp:
            gradients = self.optimizer.get_unscaled_gradients(gradients)

        # 梯度裁剪
        if self.gradient_clip_norm is not None:
            gradients, _ = tf.clip_by_global_norm(
                gradients, self.gradient_clip_norm
            )

        return gradients, tf.reduce_mean(loss) * self.accumulation_steps

    def train_epoch(self, train_dataset, epoch: int) -> dict:
        total_loss = 0.0
        total_samples = 0

        for step, (inputs, targets) in enumerate(train_dataset):
            gradients, loss = self.train_step(inputs, targets)

            # 梯度累积
            if self._gradient_buffer is None:
                self._gradient_buffer = [
                    tf.Variable(tf.zeros_like(g), trainable=False)
                    for g in gradients
                ]

            for buf, grad in zip(self._gradient_buffer, gradients):
                buf.assign_add(grad)

            self._accumulation_count += 1
            total_loss += loss.numpy() * inputs.shape[0]
            total_samples += inputs.shape[0]

            # 累积达到步数后更新参数
            if self._accumulation_count % self.accumulation_steps == 0:
                self.optimizer.apply_gradients(
                    zip(self._gradient_buffer, self.model.trainable_variables)
                )

                # 清空梯度缓冲区
                for buf in self._gradient_buffer:
                    buf.assign(tf.zeros_like(buf))

        return {"train_loss": total_loss / max(total_samples, 1)}

3.3 分布式数据并行（DDP）集成

# === PyTorch DDP ===
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_ddp_pytorch():
    """初始化 PyTorch 分布式环境"""
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def train_ddp_pytorch():
    local_rank = setup_ddp_pytorch()

    model = MyModel().cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    trainer = PyTorchTrainer(model, optimizer, use_amp=True)

    # 分布式采样器确保每个 GPU 处理不同数据
    sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    train_loader = DataLoader(
        train_dataset,
        batch_size=32,
        sampler=sampler,
        num_workers=4,
        pin_memory=True,
    )

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # 确保每轮数据打乱不同
        trainer.train_epoch(train_loader, epoch, torch.device(f"cuda:{local_rank}"))

    dist.destroy_process_group()


# === TensorFlow MirroredStrategy ===
def train_ddp_tensorflow():
    strategy = tf.distribute.MirroredStrategy()

    with strategy.scope():
        model = build_model()
        optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)

    # 分布式数据集
    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    train_dataset = strategy.experimental_distribute_dataset(
        train_dataset.batch(32).prefetch(tf.data.AUTOTUNE)
    )

    trainer = TFTrainer(model, optimizer, use_amp=True)

    for epoch in range(num_epochs):
        trainer.train_epoch(train_dataset, epoch)

四、双框架选型的工程权衡

4.1 调试体验

PyTorch 的动态图允许使用 Python 原生调试器（pdb）逐行检查张量值。TensorFlow 的 @tf.function 编译后无法直接调试，需要切换到 Eager Mode。在模型开发阶段，PyTorch 的调试效率显著更高。

4.2 生产部署

TensorFlow 的 SavedModel 格式支持跨语言部署（C++、Java、Go），TensorFlow Serving 提供了成熟的在线推理服务。PyTorch 的 TorchScript 导出仍在发展中，生产部署通常需要转换为 ONNX 格式。在部署环节，TensorFlow 的生态更成熟。

4.3 分布式训练的易用性

PyTorch DDP 需要手动管理进程启动、数据采样和梯度同步，但控制粒度更细。TensorFlow 的 Strategy API 封装程度更高，一行代码切换单卡/多卡/多机，但自定义空间受限。对于标准多卡训练，TensorFlow 更省心；对于复杂的分布式拓扑，PyTorch 更灵活。

五、总结

自定义训练循环和分布式数据并行是深度学习工程化的核心技能。两个框架的设计哲学决定了不同的实践路径：PyTorch 以动态图为核心，调试直观，控制精细，适合研究和快速迭代；TensorFlow 以静态图优化为核心，部署成熟，分布式封装度高，适合生产环境。关键选型原则：研究和原型阶段选 PyTorch，利用其调试效率和灵活性；生产部署阶段考虑 TensorFlow，利用其部署生态和分布式易用性。框架是工具，理解训练循环的每个步骤才是核心能力。