PyTorch中Module类的使用方法，模型中参数的添加与获取以及保存与载入模型

1 Module类的使用方法

1.1 add_module()方法

可以使用该方法向模型中添加层，示例如下：

self.Linear1 = nn.Linear(inputdim,hiddendim)

self.add_module("Linear1", nn.Linear(inputdim,hiddendim))

这两行代码是等价的，都可以向模型里添加两个全连接层。

1.2 获取模型中各层的信息

1.2.1 children()方法和named_children()方法

for module in model.children():
    print(module)

for name,module in model.named_children():
    print(name, module)

用这两种方法均可以得到各层信息，区别在于后者可以得到名字信息而前者不能。

1.2.2 modules()方法

用该方法可以得到整个网络的结构信息，调用方法同上。

2 模型中参数的添加与获取

2.1 为模型添加参数

2.1.1 为模型添加parameter参数

register_parameter(name,param)

2.1.2 为模型添加状态参数

register_buffer(name,tensor)

2.2 从模型中获取参数

2.2.1 用parameters()方法或named_parameters()方法获取参数

for param in model.parameters():
    print(type(param.data), param.size())

for param in model.named_parameters():
    print(type(param.data), param.size(), name)

用两种方法均可以获得模型中的参数信息，区别在于后者还可以获得参数的名字信息。

2.2.2 用state_dict()方法获取模型的全部参数

for param in model.state_dict():
    print(param, model.state_dict()[param])

'''or'''

param=model.state_dict()
print(param['conv1.weight']

用该方法可以获得模型的全部参数信息或者某个特定的参数的信息。

3 保存与载入模型

3.1 保存模型

torch.save(model.state_dict(), './model.pth')

3.2 载入模型

model.load_state_dict(torch.load('./model.pth')

model.load_state_dict(torch.load('./model.pth', map_location={'cuda:1':'cuda:0'}))

第一种是常规载入方法，第二种将模型载入到了特定的设备中。