【Mindspore】【训练过程中报错】nn.Conv2d在第二次执行时报错

问题描述：

Mindspore自定义loss时，第二次执行nn.Conv2d报错问题。且反复执行同一行代码，第一次报错，第二次不报错。由于原工程过于庞大，这里的代码进行了简化，对optimizer无需过多关注，仅需关注CustomWithLossCell的construct中nn.Conv2d计算问题。

【操作步骤&问题现象】

1、按照以下代码执行。当自定义loss时，在construct中存在两次nn.Conv2d，无论输入数据如何变化（只要shape一致），第二次的nn.Conv2d必定会报错。

import numpy as np
import mindspore as ms
from mindspore import Tensor, nn, Parameter

class CustomWithLossCell(nn.Cell):
    def __init__(self, x, w, w2):
        super(CustomWithLossCell, self).__init__()
        self.x = x
        self.w = w
        self.w2 = w2

    def construct(self):
        out = nn.Conv2d(in_channels=self.x.shape[1], out_channels=self.w.shape[0],
                        kernel_size=self.w.shape[-1], weight_init=self.w)(self.x)
        out2 = nn.Conv2d(in_channels=out.shape[1], out_channels=self.w2.shape[0],
                         kernel_size=self.w2.shape[-1], weight_init=self.w2)(out)
        return out2


x = Tensor(np.random.randn(16, 512, 4, 4), ms.float32)
w = Parameter(Tensor(np.random.randn(512, 512, 3, 3), ms.float32))
w2 = Parameter(Tensor(np.random.randn(3, 512, 1, 1), ms.float32))
w.name = "w"
w2.name = "w2"
qloss = CustomWithLossCell(x, w, w2)
optimizer = nn.Adam([w, w2], learning_rate=Tensor(0.01, ms.float32))
network = nn.TrainOneStepCell(qloss, optimizer)
network.set_train()
loss = network()

报错内容为：

RuntimeError: build/mindspore/merge/mindspore/core/ops_merge.cc:6761 Conv2dInferShape] For 'Conv2D', w_shape[0] = 512 must be equal to out_channel: 3

此时权重self.w2.shape=[3, 512, 1, 1]，输入数据out.shape=[16, 512, 4, 4]，正确输出的shape应为[16, 3, 4, 4]。报错内容无法理解，weight.shape[0]应该就是3，却检测出来是512。也无法寻找到ops_merge.cc文件在电脑中的具体位置。

2、若在第二个nn.Conv2d处添加断点，在Console中运行out2 = nn.Conv2d(in_channels=out.shape[1], out_channels=self.w2.shape[0], kernel_size=self.w2.shape[-1], weight_init=self.w2)(out)，第一次会报同样的错，第二次再次运行该行，则不会报错，且输出的结果正确。尽管第二次可以成功运行，但由于计算了两次，梯度会丢失，梯度全为0。

【截图信息】

源代码：

 

报错：

终端第二次运行： 

解答：

这个问题是因为在TrainOneStepCell构图时，因为上述的CustomWithLossCell是在construct里调用nn.Conv2d的，Conv2d的权重Parameter名字无法区分导致第二次调用的Conv2d的权重使用了第一次的缓存。

需要把nn.Conv2d初始化放到__init__里，改成如下的方式：

class CustomWithLossCellV2(nn.Cell):

    def __init__(self, x, w, w2):

        super(CustomWithLossCellV2, self).__init__()

        self.x = x

        self.w = w

        self.w2 = w2

        self.conv1 = nn.Conv2d(in_channels=self.x.shape[1], out_channels=self.w.shape[0],

                               kernel_size=self.w.shape[-1])

        self.conv2 = nn.Conv2d(in_channels=self.w.shape[0], out_channels=self.w2.shape[0],

                               kernel_size=self.w2.shape[-1])



    def construct(self):

        out = self.conv1(self.x)

        out2 = self.conv2(out)

        return out2

这种方式，在Cell里能看到Conv2d的权重，所以可以给它们赋以不同的名字。