微分接口适配

适配MindSpore框架原生微分接口

由于MindSpore的自动微分采用函数式表达,和PyTorch的微分接口存在差异,目前需要用户手动适配训练部分的少量代码,即将PyTorch版本的训练流程代码转换为MindSpore的函数式编程写法,从而使能MindSpore动静统一、自动并行等竞争力功能。详细内容可参考MindSpore使用文档。以下示例展示了如何将PyTorch训练流程转换为MindSpore函数式训练流程:

迁移前网络表达:

net = LeNet().to(config_args.device)
optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0005)
net.train()

# 数据迭代训练
for i in range(epochs):
    for X, y in train_data:
        X, y = X.to(config_args.device), y.to(config_args.device)
        out = net(X)
        loss = criterion(out, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print("------>epoch:{}, loss:{:.6f}".format(i, loss))

替换为Mindspore函数式迭代训练表达,其中前向过程通常包含了模型网络接口调用以及损失函数调用,反向求导过程包含了反向梯度接口调用以及优化器接口调用部分,此外,MindSpore不需要调用loss.backward()以及optimizer.zero_grad(),具体示例如下:

import mindtorch.torch as torch
import mindspore as ms

net = LeNet().to(config_args.device)
optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0005)

# 定义前向过程
def forward_fn(data, label):
    logits = net(data)
    loss = criterion(logits, label)
    return loss, logits

# 反向梯度定义
grad_fn = ms.ops.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)

# 单步训练定义
def train_step(data, label):
    (loss, _), grads = grad_fn(data, label)
    optimizer(grads)
    return loss

net.train()
# 数据迭代训练
for i in range(epochs):
    for X, y in train_data:
        X, y = X.to(config_args.device), y.to(config_args.device)
        res = train_step(X, y)
        print("------>epoch:{}, loss:{:.6f}".format(i, res.numpy()))