使用Mindspore模型训练时出现梯度为0现象

1 系统环境

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

2 报错信息

2.1问题描述

使用Mindspore模型训练时出现梯度为0现象

2.2 脚本代码

def forward_fn(inputs, targets):
    logits = model(inputs)
    loss = criterion(logits, targets)
    return loss, logits

# get grad function
grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)

# define train step function
def train_step(inputs, targets):
    (loss, logits), grads = grad_fn(inputs, targets) # get values and gradients
    print('grads=',grads)
    optimizer(grads) # update gradient
    return loss, logits

3 根因分析

分析代码,梯度是在value_and_grad中获取的,看看官网的介绍,第三个位置的weight需要net.trainable_params() 获取,代码里的net是model
cke_21565.png

4 解决方案

修改

optimizer.parameters为 model.trainable_params() 
# get grad function 
grad_fn = ms.value_and_grad(forward_fn, None, model.trainable_params(), has_aux=True)