MindSpore不能像torch的param.grad直接获取梯度问题及解决

1.系统环境

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

2. 问题描述

torch代码,对于parameter变量,例如X,可以直接X.grad获取梯度,torch封装在内部逻辑中,不需要额外的代码计算梯度,MindSpore不能像torch的param.grad直接获取梯度。

import torch  
x = torch.tensor(2., requires_grad=True)  
y = torch.tensor(3., requires_grad=True)  
z = x * x * y  
z.backward()  
print(x.grad, y.grad)  
# Out:  
# tensor(12.) tensor(4.)  
    
x = torch.tensor(2.).requires_grad_()  
y = torch.tensor(3.).requires_grad_()  
z = x * x * y  
grad_x = torch.autograd.grad(outputs=z, inputs=x)  
print(grad_x[0])

3. 解决方案

PyTorch:使用torch.autograd.backward计算给定Tensor关于叶子节点的梯度总和,反向传播计算Tensor的梯度时,只计算requires_grad=True的叶子节点的梯度。使用torch.autograd.grad计算并返回输出关于输入的梯度总和,如果only_inputs为True,仅返回与指定输入相关的梯度列表。

MindSpore:计算梯度,当grad_position设置为int或者tuple int类型,将会计算对应输入位置的梯度。如果设置了weights,将会计算网络的变量的参数。当has_aux设置为True时,只有fn的第一个输出参与梯度计算,此时fn至少具备两个输出。

import numpy as np  
import mindspore.nn as nn  
import mindspore as ms  
from mindspore import ops  
    
# In MindSpore:  
class Net(nn.Cell):  
    def __init__(self):  
        super(Net, self).__init__()  
        self.matmul = ops.MatMul()  
        self.z = ms.Parameter(ms.Tensor(np.array([1.0], np.float32)), name='z')  
    def construct(self, x, y):  
        x = x * self.z  
        out = self.matmul(x, y)  
        return out  
    
class GradNetWrtX(nn.Cell):  
    def __init__(self, net):  
        super(GradNetWrtX, self).__init__()  
        self.net = net  
    def construct(self, x, y):  
        gradient_function = ms.grad(self.net)  
        return gradient_function(x, y)  
    
x = ms.Tensor([[0.5, 0.6, 0.4], [1.2, 1.3, 1.1]], dtype=ms.float32)  
y = ms.Tensor([[0.01, 0.3, 1.1], [0.1, 0.2, 1.3], [2.1, 1.2, 3.3]], dtype=ms.float32)  
output = GradNetWrtX(Net())(x, y)  
print(output)