关于pytorch中module.apply算子在MindSpore下代替实现的方法

转载地址：https://bbs.huaweicloud.com/forum/thread-137700-1-1.html

作者：Topia_

我们在使用pytorch时，有时会用到一个非常好用的函数：module.apply()。但是在MindSpore里并没有直接对应的函数，下面我们来分析替代实现这个算子的方法。

pytorch的任何网络net，都是torch.nn.Module的子类,都算是module，也就是模块。pytorch中的model.apply(fn)会递归地将函数fn应用到父模块的每个子模块submodule，也包括model这个父模块自身。

例如：

import torch.nn as nn
net = nn.Sequential(nn.Conv2d(3,4,5),nn.BatchNorm2d(2))

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find("BatchNorm2d") != -1:
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
        torch.nn.init.constant_(m.bias.data, 0.0)

net.apply(weights_init_normal)

这样，只用apply这个函数，就可以一次性实现对所有子模块的初始化权重赋值。

那么，对于MindSpore使用的nn.Cell,虽然没有“cell.apply”这样的算子，但是实际上nn.Cell提供了更为简单的操作方法，即cells_and_names()，我们先看一下这个属性具体怎么使用：

n = Net()
names = []
for m in n.cells_and_names():
    if m[0]:
        names.append(m[0])

以上是mindspore官网给出的例程，实际上，n.cells_and_names就是一个迭代器，用变量m就可以遍历网络Net()的所有子模块。其中m[0]表示模块名，m[1]表示模块自身.
不同的是，如果用SequentialCell直接将子模块连接起来，m[0]则为0, 0.0, 0.1, 0.2, 1, 1.1, ...,例如：

import mindspore.nn as nn
net = nn.SequentialCell(nn.SequentialCell(nn.Conv2d(3,4,5),nn.ReLU()), 
                        nn.ReLU())
for m in net.cells_and_names():
    print(m[0])
    print(m[1])

输出：

SequentialCell<
  (0): SequentialCell<
    (0): Conv2d
    (1): ReLU<>
    >
  (1): ReLU<>
  >
0
SequentialCell<
  (0): Conv2d
  (1): ReLU<>
  >
0.0
Conv2d
0.1
ReLU<>
1
ReLU<>

如果定义一个类再做处理：

import mindspore.nn as nn
class test(nn.Cell):
    def __init__(self):
        super(test, self).__init__()
        self.conv2d = nn.Conv2d(3,4,5)
        self.relu = nn.ReLU()
    def construct(self,x):
        x = self.conv2d(x)
        return self.relu(x)

net = test()
for m in net.cells_and_names():
    if m[0] == 'conv2d':
        print("Include conv2d module.")
    elif m[0] == 'relu':
        print("Include relu module.")

输出：

Include conv2d module.
Include relu module.

掌握了cells_and_names的使用方法后，就可以自由对nn.Cell内的子模块进行遍历，然后根据需求定义自己的操作了。

总结：pytorch下需要先定义函数，再用apply调用函数才能进行处理；MindSpore只需要定义一次函数即可完成操作，更直观更简洁。