pytorch和paddle的存储模型变量state_dict命名规则分析

一、Pytorch存储模型变量命名分析

在pytorch中，存储变量的名称就在def init(self)中定义，名字就是self中的定义名称。若在类中还调用了其他的类，那么名称则为实例化的变量名称。

典型示例如下：

import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
 
 
class test(nn.Module):
    '''
    定义子类
    '''
    def __init__(self):
        super(test, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)


class TheModelClass(nn.Module):
    '''
    定义测试类
    '''
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        self.test= test()
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


def main():
    '''
    主函数
    '''
    # 建立模型
    model = TheModelClass()

    # 建立优化器
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 输出模型变量字典
    print('Model.state_dict:')
    for param_tensor in model.state_dict():
        # 打印 key value字典
        print(param_tensor, '\t', model.state_dict()[param_tensor].size())

    # 输出优化器变量字典
    print('Optimizer,s state_dict:')
    for var_name in optimizer.state_dict():
        print(var_name, '\t', optimizer.state_dict()[var_name])


if __name__ == '__main__':
    '''
    程序入口
    '''
    main()

输出结果如下：

Model.state_dict:
conv1.weight     torch.Size([6, 3, 5, 5])
conv1.bias       torch.Size([6])
conv2.weight     torch.Size([16, 6, 5, 5])
conv2.bias       torch.Size([16])
fc1.weight       torch.Size([120, 400])
fc1.bias         torch.Size([120])
fc2.weight       torch.Size([84, 120])
fc2.bias         torch.Size([84])
fc3.weight       torch.Size([10, 84])
fc3.bias         torch.Size([10])
test.conv1.weight        torch.Size([6, 3, 5, 5])
test.conv1.bias          torch.Size([6])
test.conv2.weight        torch.Size([16, 6, 5, 5])
test.conv2.bias          torch.Size([16])
Optimizer,s state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]}]

二、PaddlePaddle存储模型变量命名分析

典型代码如下：

import paddle.nn as nn
import paddle.optimizer as optim
import paddle.nn.functional as F
 
 
class test(nn.Layer):
    '''
    定义子类
    '''
    def __init__(self):
        super(test, self).__init__()
        self.conv1 = nn.Conv2D(3, 6, 5)
        self.pool = nn.MaxPool2D(2, 2)
        self.conv2 = nn.Conv2D(6, 16, 5)


class TheModelClass(nn.Layer):
    '''
    定义测试类
    '''
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2D(3, 6, 5)
        self.pool = nn.MaxPool2D(2, 2)
        self.conv2 = nn.Conv2D(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        self.test= test()
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


def main():
    '''
    主函数
    '''
    # 建立模型
    model = TheModelClass()

    # 建立优化器
    optimizer = optim.SGD(parameters = model.parameters(), learning_rate=0.001, weight_decay=0.9)

    # 输出模型变量字典
    print('Model.state_dict:')
    for param_tensor in model.state_dict().keys():
        # 打印 key value字典
        print(param_tensor, '\t', model.state_dict()[param_tensor].shape)


if __name__ == '__main__':
    '''
    程序入口
    '''
    main()

其输出如下所示：

Model.state_dict:
conv1.weight     [6, 3, 5, 5]
conv1.bias       [6]
conv2.weight     [16, 6, 5, 5]
conv2.bias       [16]
fc1.weight       [400, 120]
fc1.bias         [120]
fc2.weight       [120, 84]
fc2.bias         [84]
fc3.weight       [84, 10]
fc3.bias         [10]
test.conv1.weight        [6, 3, 5, 5]
test.conv1.bias          [6]
test.conv2.weight        [16, 6, 5, 5]
test.conv2.bias          [16]

通过对比发现，在命名规则上pytorch和paddlepaddle是一样的。只不过对于fc层来说，它的weight的形状是相互转置的关系。

三、Pytorch和Paddle相互转化

通过上面的分析我们知道，pytorch和paddle的模型变量命名规则是完全一样的。那么对于训练好的pytorch或paddle模型，我们就可以基于上述原则进行互转。在互换时注意fc层，对于fc层的变量需要做转置处理。