使用MindSpore实现pytorch中的前反向传播

经验分享 Tech Blogs 功能调试-Function Debugging
1

1 系统环境

硬件环境(Ascend/GPU/CPU): Ascend/GPU/CPU
MindSpore版本: mindspore=2.2.0
执行模式（PyNative/ Graph）:不限
Python版本: Python=3.9
操作系统平台: 不限

2 报错信息

2.1 问题描述

如何使用MindSpore实现pytorch中的前向传播和反向传播方式

2.2 脚本信息

output = rnn(b_x)  # rnn output
loss = loss_func(output, b_y)  # cross entropy loss
optimizer.zero_grad()  # clear gradients for this training step
loss.backward()  # backpropagation, compute gradients
optimizer.step()  # apply gradients

3 根因分析

mindspore下：
在模型训练中，一个完整的训练过程（step）需要实现以下三步：

  1. 正向计算 ：模型预测结果（logits），并与正确标签（label）求预测损失（loss）。
  2. 反向传播 ：利用自动微分机制，自动求模型参数（parameters）对于loss的梯度（gradients）。
  3. 参数优化 ：将梯度更新到参数上。
    MindSpore使用函数式自动微分机制，因此针对上述步骤需要实现：
  4. 定义正向计算函数。
  5. 使用value_and_grad通过函数变换获得梯度计算函数。
  6. 定义训练函数，使用set_train设置为训练模式，执行正向计算、反向传播和参数优化。

4 解决方案

以下用一个简单的网络来对比torch和mindspore 前向传播和反向传播的差别

import torch  
import torch.nn as nn  
import torch.optim as optim  
    
    
# 定义一个简单的模型  
class Model(nn.Module):  
    def __init__(self):  
        super(Model, self).__init__()  
        self.weight = nn.Parameter(torch.randn(1))  
        self.bias = nn.Parameter(torch.randn(2))  
    
    
    def forward(self, x):  
        y = self.weight*x +self.bias  
        return y  
    
    
# 实例化模型  
model = Model()  
    
# 随机的输入和目标  
inputs = torch.tensor([1.,2.], dtype=torch.float32, requires_grad=True)  
    
targets = torch.tensor([101., 202.])  
    
# 定义损失函数  
loss_fn = nn.MSELoss()  
    
# 定义优化器，这里以SGD为例  
optimizer = optim.SGD(model.parameters(), lr=0.01)  
    
    
for i in range(1000):  
    print(f"-----------------------loop:{i}---------------------")  
    # 前向传播  
    outputs = model(inputs)  
    loss = loss_fn(outputs, targets)  
    
    print("loss:", loss)  
    if loss.detach().numpy() <0.001:  
        break  
    # 反向传播  
    loss.backward()  
    print("inputs.grad:", inputs.grad)  
    
    print("inputs before optimizer.step:", model.parameters())  
    for i in  model.parameters():  
        print(i)  
    # 梯度更新  
    optimizer.step()  
    print("inputs after optimizer.step:", model.parameters())  
    for i in  model.parameters():  
        print(i)  
    optimizer.zero_grad()

运行结果：

-----------------------loop:0---------------------  
loss: tensor(25024.0586, grad_fn=<MseLossBackward0>)  
inputs.grad: tensor([-107.0528, -217.2061])  
inputs before optimizer.step: <generator object Module.parameters at 0x0000029FE4BF6900>  
Parameter containing:  
tensor([1.0824], requires_grad=True)  
Parameter containing:  
tensor([ 1.0168, -0.8306], requires_grad=True)  
inputs after optimizer.step: <generator object Module.parameters at 0x0000029FE4BF6900>  
Parameter containing:  
tensor([6.0848], requires_grad=True)  
Parameter containing:  
tensor([2.0058, 1.1760], requires_grad=True)  
-----------------------loop:1---------------------  
loss: tensor(22111.3359, grad_fn=<MseLossBackward0>)  
inputs.grad: tensor([ -672.3835, -1365.1215])  
inputs before optimizer.step: <generator object Module.parameters at 0x0000029FE4BF6900>  
Parameter containing:  
tensor([6.0848], requires_grad=True)  
Parameter containing:  
tensor([2.0058, 1.1760], requires_grad=True)  
inputs after optimizer.step: <generator object Module.parameters at 0x0000029FE4BF6900>  
Parameter containing:  
tensor([10.7869], requires_grad=True)  
Parameter containing:  
tensor([2.9349, 3.0626], requires_grad=True)  
。。。。。。  
  
-----------------------loop:333---------------------  
loss: tensor(0.0010, grad_fn=<MseLossBackward0>)  
inputs.grad: tensor([ -61202.0000, -142309.9844])  
inputs before optimizer.step: <generator object Module.parameters at 0x0000029FE4BF6900>  
Parameter containing:  
tensor([84.4544], requires_grad=True)  
Parameter containing:  
tensor([16.5859, 33.0709], requires_grad=True)  
inputs after optimizer.step: <generator object Module.parameters at 0x0000029FE4BF6900>  
Parameter containing:  
tensor([84.4544], requires_grad=True)  
Parameter containing:  
tensor([16.5855, 33.0712], requires_grad=True)  
-----------------------loop:334---------------------  
loss: tensor(0.0010, grad_fn=<MseLossBackward0>)  

from mindspore import nn,Parameter  
import mindspore as ms  
from mindspore import Tensor, ops  
import numpy as np  
# 定义一个简单的模型  
class Model(nn.Cell):  
    def __init__(self):  
        super(Model, self).__init__()  
        self.weight = Parameter(Tensor(np.random.randn(1),dtype=ms.float32))  
        self.bias = Parameter(Tensor(np.random.randn(2),dtype=ms.float32))  
    
    
    def construct(self, x):  
        y = self.weight*x +self.bias  
        return y  
# 实例化模型  
model = Model()  
    
# 随机的输入和目标  
inputs = Tensor([1.,2.],dtype=ms.float32)  
    
targets = Tensor([101., 202.],dtype=ms.float32)  
optimizer = nn.SGD(model.trainable_params(), learning_rate=0.01)  
# 定义损失函数  
loss_fn = nn.MSELoss()  
    
#由于需要使用函数式自动微分，需要将神经网络和损失函数的调用封装为一个前向计算函数  
# Define forward function  
def forward_fn(inputs, target):  
    result = model(inputs)  
    loss = loss_fn(result, target)  
    return loss  
#使用value_and_grad通过函数变换获得梯度计算函数  
grad_fn = ms.value_and_grad(forward_fn, None, weights=model.trainable_params())  
model.set_train()  
for i in range(1000):  
    print(f"-----------------------loop:{i}---------------------")  
    # 前向传播  
    outputs = model(inputs)  
    # 反向传播  
    loss, grads = grad_fn(inputs, targets)  
    
    print("loss:", loss)  
    if loss.numpy() <0.001:  
        break  
    print("inputs before optimizer:", model.get_parameters())  
    for i in model.get_parameters():  
        print(i.value())  
    # 梯度更新  
    optimizer(grads)  
    print("inputs after optimizer:", model.get_parameters())  
    for i in  model.get_parameters():  
        print(i.value())  

-----------------------loop:0---------------------  
loss: 25290.72  
inputs before optimizer: <generator object Cell.get_parameters at 0x000002272CBE23C0>  
[-0.31287298]  
[-0.07084849  1.8702196 ]  
inputs after optimizer: <generator object Cell.get_parameters at 0x000002272DCF5200>  
[4.716075]  
[0.94298875 3.8777747 ]  
-----------------------loop:1---------------------  
loss: 22346.92  
inputs before optimizer: <generator object Cell.get_parameters at 0x000002272DCF5200>  
[4.716075]  
[0.94298875 3.8777747 ]  
inputs after optimizer: <generator object Cell.get_parameters at 0x000002272DCF5200>  
[9.443286]  
[1.8963981 5.7646756]  
。。。。。  
  
  
-----------------------loop:298---------------------  
loss: 0.0010135988  
inputs before optimizer: <generator object Cell.get_parameters at 0x000002272DCF5200>  
[83.50287]  
[17.456827 35.014324]  
inputs after optimizer: <generator object Cell.get_parameters at 0x000002272DCF5200>  
[83.50287]  
[17.45723  35.014122]  
-----------------------loop:299---------------------  
loss: 0.0009934219

具体差别查看可以看代码，都有相关注释。

  • 网络定义：在网络定义中，一般会定义出需要的前向网络，损失函数和优化器。在Net()中定义前向网络，PyTorch的网络继承nn.Module；类似地，MindSpore的网络继承nn.Cell。在MindSpore中，损失函数和优化器除了使用MindSpore中提供的外，用户还可以使用自定义的优化器。可参考模型模块自定义。可以使用functional/nn等接口拼接需要的前向网络、损失函数和优化器，详细接口用法可参照接口对比。
  • 正向计算：运行实例化后的网络，可以得到logit，将logit和target作为输入计算loss。需要注意的是，如果正向计算的函数有多个输出，在反向计算时需要注意多个输出对于计算结果的影响。
  • 反向计算：得到loss后，我们可以进行反向计算。在PyTorch中可使用loss.backward()计算梯度，在MindSpore中，先用mindspore.grad()定义出反向传播方程net_backward，再将输入传入net_backward中，即可计算梯度。如果正向计算的函数有多个输出，在反向计算时，可将has_aux设置为True，即可保证只有第一个输出参与求导，其它输出值将直接返回。对于反向计算中接口用法区别详见自动微分对比。
  • 梯度更新：将计算后的梯度更新到网络的Parameters中。在PyTorch中使用optim.step()；在MindSpore中，将Parameter的梯度传入定义好的optim中，即可完成梯度更新。
【MindSpore报错解决地图】常见报错问题和解决方案（持续更新）