Pytorch基础——从nn.module转写成.py脚本（一）

简介

起因还是在做AutoML，中间有小伙伴需要改写模型中channel的数量，希望能将模型channel变成一个变量，这需要改动torch定义的model类的脚本文件，所以想能不能根据传入的model，自动生成.py的脚本。

代码

为了测试代码，我们先写一个model.py文件生成一个简单地模型。

# model.py

import torch
import torch.nn as nn


class Model(nn.Module):
    
    def __init__(self,):
        super(Model, self).__init__()

        self.input = nn.Sequential(
            nn.Conv2d(1,16,3,1),
            nn.BatchNorm2d(16),
        )

        self.convs = nn.Sequential(
            nn.Conv2d(16,100,3),
            nn.ReLU(),
            nn.Conv2d(100,100,3),
            nn.BatchNorm2d(16),
            nn.Conv2d(100,100,3),
            nn.MaxPool2d(2),
            nn.Conv2d(100,100,3),            
        )
        self.out = nn.Linear(100,10)
        self.initialize()
    
    def initialize(self,):
        pass

    def forward(self, x):

        x = self.input(x)
        x2 = self.convs(x)
        x = x + x2
        x = self.out(x)

        return x

if __name__ == "__main__":
    model = Model()
    print(model.named_modules)

上面这部分就不多说了，下面是正文。新建一个文件，我们暂且命名为autoCode.py

from model import Model

model = Model()

先导入刚才构建的模型类，实例一个model出来。
为了简单起见，我们手工定义一下文件头部的内容

global tab
tab  = '    '
model_name = str(model.__class__.__name__)
file = './script/'+ model_name +'.py'
file_head = '# -*-coding:utf-8-*-\n\nimport torch\nimport torch.nn as nn\n\nclass '\
                + str(model.__class__.__name__)+'(nn.Module):\n'+tab+'def __init__(self,):\n'+tab+tab\
                + 'super(' + model_name + ', self).__init__()\n'

因为pytorch中是通过_module子模块对nn.module进行层层封装来得到最终的模型的，我们对它写一个递归调用，来获取每一个封装里的子操作。需要说明的是，._module是一个orderedDict类型变量，因此可以用.items()来得到它的key和value，其中key就是这个操作层的名字，value就是具体的操作，如value: Conv2d(…)，我们要生成的代码就是value部分的内容。不过在torch定义中这些操作都在torch.nn这个类下，因此需要在value前面加上nn.。

def nn_replace(model):
    global tab
    script = ''
    tmp_script = ''
    if len(model._modules) == 0:
        # name, m = model.named_modules()
        # print('module: nn.{},'.format(str(model)))
        script += 'nn.{},\n'.format(str(model))
        return script, tmp_script
        # model.pop(name)
    else:
        for n,m in model._modules.items():
            name = n
            s,_ = nn_replace(m)
            script += s
            try:
                int(n)
                continue
            except ValueError:
                pass
            # print('name: {}'.format(name))
            script = 'self.' + name + ' = nn.Sequential(' + script + ')\n'
            tmp_script += tab + tab + script
            script = ''
    return script, tmp_script

最后我们调用上述函数，返回script，并添加到head_script里，一起写到文件里看看

_, script = nn_replace(model)
print(script)

file_head += script
with open(file,'w') as f:
    f.write(file_head)

out: self.input = nn.Sequential(nn.Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1)),
nn.BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
)
self.convs = nn.Sequential(nn.Conv2d(16, 100, kernel_size=(3, 3), stride=(1, 1)),
nn.ReLU(),
nn.Conv2d(100, 100, kernel_size=(3, 3), stride=(1, 1)),
nn.BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
nn.Conv2d(100, 100, kernel_size=(3, 3), stride=(1, 1)),
nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False),
nn.Conv2d(100, 100, kernel_size=(3, 3), stride=(1, 1)),
)
self.out = nn.Sequential(nn.Linear(in_features=100, out_features=10, bias=True),
)

结果

生成的代码保存在./script/目录下，文件名与之前实例化的模型的l类名称相同，这里是Model.py

# -*-coding:utf-8-*-

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self,):
        super(Model, self).__init__()
        self.input = nn.Sequential(nn.Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1)),
nn.BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
)
        self.convs = nn.Sequential(nn.Conv2d(16, 100, kernel_size=(3, 3), stride=(1, 1)),
nn.ReLU(),
nn.Conv2d(100, 100, kernel_size=(3, 3), stride=(1, 1)),
nn.BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
nn.Conv2d(100, 100, kernel_size=(3, 3), stride=(1, 1)),
nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False),
nn.Conv2d(100, 100, kernel_size=(3, 3), stride=(1, 1)),
)
        self.out = nn.Sequential(nn.Linear(in_features=100, out_features=10, bias=True),
)

可以看到对的不齐，但并不关心。括号内部的空格都会被忽略。试验一下这个代码能不能生成模型，在下面添加以下内容：

if __name__ == "__main__":
    model = Model()
    print(model)

结果如下：

Model(
(input): Sequential(
(0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1))
(1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(convs): Sequential(
(0): Conv2d(16, 100, kernel_size=(3, 3), stride=(1, 1))
(1): ReLU()
(2): Conv2d(100, 100, kernel_size=(3, 3), stride=(1, 1))
(3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(4): Conv2d(100, 100, kernel_size=(3, 3), stride=(1, 1))
(5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(6): Conv2d(100, 100, kernel_size=(3, 3), stride=(1, 1))
)
(out): Sequential(
(0): Linear(in_features=100, out_features=10, bias=True)
)
)

it works.
这个小脚本只解决了部分简单地网络模型到代码生成的问题，对于其中的channel还没有自动改成变量，后续组里小伙伴如果想到了再更新吧。