【MindSpore进阶】抛弃Model.train，手把手教你写自定义训练流

昇思MindSpore 2025-12-05 11:31:16

MindSpore 的 Model.train接口非常方便，像 Keras 一样封装了所有细节。但在科研或复杂的工程落地场景中（例如 GAN 网络、强化学习、或者需要魔改梯度更新策略时），高层接口往往显得不够灵活。

很多从 PyTorch 转到 MindSpore 的同学都在问：“怎么写一个纯手动的 for 循环训练代码？”

在 MindSpore 2.x 版本推荐的函数式编程范式下，自定义训练流其实非常优雅。今天我们就来拆解如何在昇腾上实现一套灵活的训练循环。

核心思想：函数变换

PyTorch 是基于对象的（optimizer.step()），而 MindSpore 推荐基于函数的（Function Transformation）。我们需要构建一个前向计算函数，然后通过 ops.value_and_grad自动生成梯度计算函数。

实战代码模板

下面是一个完整的、可运行的自定义训练模板。

1. 准备环境与网络

import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
import numpy as np

ms.set_context(mode=ms.GRAPH_MODE, device_target="Ascend")

一个简单的网络

class Net(nn.Cell):
    def init(self):
        super(Net, self).init()
        self.fc = nn.Dense(10, 1)
    def construct(self, x):
        return self.fc(x)

模拟数据

data = ms.Tensor(np.random.randn(32, 10).astype(np.float32))
label = ms.Tensor(np.random.randn(32, 1).astype(np.float32))

net = Net()
loss_fn = nn.MSELoss()
optimizer = nn.Momentum(net.trainable_params(), learning_rate=0.01, momentum=0.9)

2. 定义Forward函数（核心）

这里我们不定义一个 TrainOneStepCell类，而是直接写 Python 函数。注意，这个函数包含了 Loss 的计算逻辑。

def forward_fn(data, label):
    logits = net(data)
    loss = loss_fn(logits, label)
    return loss, logits

3. 生成梯度计算函数（魔法所在）

使用 ops.value_and_grad。

grad_position=None：表示不对输入求导。
weights=optimizer.parameters：表示对网络权重求导。
has_aux=True：因为 forward_fn返回了 (loss, logits)两个值，我们需要告诉它第一个是 Loss（用于求导），第二个是辅助数据。

# 自动微分
grad_fn = ops.value_and_grad(forward_fn, 
                             grad_position=None, 
                             weights=optimizer.parameters, 
                             has_aux=True)

4. 构造训练Step

这一步我们需要加上 @ms.jit装饰器。这是为了将这个函数编译成静态图下沉到昇腾NPU上执行，否则它会以解释模式运行，速度很慢。

@ms.jit
def train_step(data, label):
    # 1. 计算梯度和Loss
    (loss, _), grads = grad_fn(data, label)
  
    # 2. (可选) 在这里可以对 grads 做任何你想做的骚操作
    # 例如：梯度累积、梯度裁剪、梯度加噪...
    # grads = ops.clip_by_global_norm(grads, 1.0)
  
    # 3. 优化器更新参数
    optimizer(grads)
  
    return loss

5. 训练循环

现在，你可以像写 PyTorch 一样写循环了：

epochs = 5
print("Start Training...")

for epoch in range(epochs):
    # 假设这里有一个 dataloader
    # for d, l in dataloader:
  
    # 这里直接用模拟数据演示
    loss = train_step(data, label)
  
    print(f"Epoch: {epoch+1}, Loss: {loss}")

print("Done!")

为什么这样做？

这种方式对比 Model.train有极大的优势：

透明：你可以清楚地看到梯度产生的瞬间，并有机会拦截它。
灵活：实现对抗生成网络（GAN）时，需要交替训练 Generator 和 Discriminator，这种模式只需要写两个 grad_fn和 train_step并在循环中交替调用即可。
高性能：虽然写法像 PyTorch，但 @ms.jit保证了它在昇腾芯片上依然是整图下沉执行的，并没有性能损失。