1. 重要类

• container.py
• nn.sequential
• nn.modulelist
• save_state_dict

2. 保存模型

pytorch官网教程

3. 代码测试

比较急，后续完善

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :ToTest01.py
# @Time      :2024/11/24 10:37
# @Author    :Jason Zhang
import torch
from torch import nn
from torch.nn import Moduleclass MyModel(nn.Module):def __init__(self):super(MyModel, self).__init__()self.linear1 = nn.Linear(2, 3)self.linear2 = nn.Linear(3, 4)self.batch_norm4 = nn.BatchNorm2d(4)def forward(self, x):x = self.linear1(x)x = self.linear2(x)return xif __name__ == "__main__":run_code = 0input_x = torch.randn((1, 2))test_model = MyModel()y = test_model(input_x)model_modules = test_model._modulesprint(f"*"*50)print(f"model_modules=\n{model_modules}")print(f"*"*50)linear1 = model_modules['linear1']print(f"*"*50)print(f"linear1={linear1}")print(f"*"*50)print(f"linear1.weight=\n{linear1.weight}")print(f"*"*50)print(f"linear1.weight.dtype={linear1.weight.dtype}")print(f"*"*50)test_model.to(torch.double)print(f"linear1.weight.dtype={linear1.weight.dtype}")print(f"*"*50)test_model.to(torch.float32)print(f"linear1.weight.dtype={linear1.weight.dtype}")print(f"*"*50)model_parameters = test_model._parametersprint(f"model_parameters={model_parameters}")print(f"*"*50)model_buffers = test_model._buffersprint(f"model_buffer={model_buffers}")print(f"*"*50)model_state_dict = test_model.state_dict()print(f"model_state_dict=\n{model_state_dict}")print(f"*"*50)model_state_dict_linear2 = test_model.state_dict()['linear2.weight']print(f"model_state_dict_linear2=\n{model_state_dict_linear2}")print(f"*"*50)model_named_para =list(test_model.named_parameters())print(f"model_named_para=\n{model_named_para}")print(f"*"*50)model_named_modules =list(test_model.named_modules())print(f"model_named_modules=\n{model_named_modules}")print(f"*"*50)model_named_buffers =list(test_model.named_buffers())print(f"model_named_buffers=\n{model_named_buffers}")print(f"*"*50)model_named_children =list(test_model.named_children())print(f"model_named_children=\n{model_named_children}")

• 结果：

**************************************************
model_modules=
OrderedDict([('linear1', Linear(in_features=2, out_features=3, bias=True)), ('linear2', Linear(in_features=3, out_features=4, bias=True)), ('batch_norm4', BatchNorm2d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))])
**************************************************
**************************************************
linear1=Linear(in_features=2, out_features=3, bias=True)
**************************************************
linear1.weight=
Parameter containing:
tensor([[-0.5518,  0.0687],[-0.7013,  0.4869],[-0.1157, -0.1287]], requires_grad=True)
**************************************************
linear1.weight.dtype=torch.float32
**************************************************
linear1.weight.dtype=torch.float64
**************************************************
linear1.weight.dtype=torch.float32
**************************************************
model_parameters=OrderedDict()
**************************************************
model_buffer=OrderedDict()
**************************************************
model_state_dict=
OrderedDict([('linear1.weight', tensor([[-0.5518,  0.0687],[-0.7013,  0.4869],[-0.1157, -0.1287]])), ('linear1.bias', tensor([-0.2915, -0.4807,  0.0071])), ('linear2.weight', tensor([[ 0.4185,  0.1556,  0.1371],[ 0.4751,  0.2029, -0.0679],[ 0.1264, -0.0288, -0.3661],[ 0.4423, -0.5370,  0.3930]])), ('linear2.bias', tensor([ 0.2746, -0.1798,  0.0218,  0.5465])), ('batch_norm4.weight', tensor([1., 1., 1., 1.])), ('batch_norm4.bias', tensor([0., 0., 0., 0.])), ('batch_norm4.running_mean', tensor([0., 0., 0., 0.])), ('batch_norm4.running_var', tensor([1., 1., 1., 1.])), ('batch_norm4.num_batches_tracked', tensor(0))])
**************************************************
model_state_dict_linear2=
tensor([[ 0.4185,  0.1556,  0.1371],[ 0.4751,  0.2029, -0.0679],[ 0.1264, -0.0288, -0.3661],[ 0.4423, -0.5370,  0.3930]])
**************************************************
model_named_para=
[('linear1.weight', Parameter containing:
tensor([[-0.5518,  0.0687],[-0.7013,  0.4869],[-0.1157, -0.1287]], requires_grad=True)), ('linear1.bias', Parameter containing:
tensor([-0.2915, -0.4807,  0.0071], requires_grad=True)), ('linear2.weight', Parameter containing:
tensor([[ 0.4185,  0.1556,  0.1371],[ 0.4751,  0.2029, -0.0679],[ 0.1264, -0.0288, -0.3661],[ 0.4423, -0.5370,  0.3930]], requires_grad=True)), ('linear2.bias', Parameter containing:
tensor([ 0.2746, -0.1798,  0.0218,  0.5465], requires_grad=True)), ('batch_norm4.weight', Parameter containing:
tensor([1., 1., 1., 1.], requires_grad=True)), ('batch_norm4.bias', Parameter containing:
tensor([0., 0., 0., 0.], requires_grad=True))]
**************************************************
model_named_modules=
[('', MyModel((linear1): Linear(in_features=2, out_features=3, bias=True)(linear2): Linear(in_features=3, out_features=4, bias=True)(batch_norm4): BatchNorm2d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)), ('linear1', Linear(in_features=2, out_features=3, bias=True)), ('linear2', Linear(in_features=3, out_features=4, bias=True)), ('batch_norm4', BatchNorm2d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))]
**************************************************
model_named_buffers=
[('batch_norm4.running_mean', tensor([0., 0., 0., 0.])), ('batch_norm4.running_var', tensor([1., 1., 1., 1.])), ('batch_norm4.num_batches_tracked', tensor(0))]
**************************************************
model_named_children=
[('linear1', Linear(in_features=2, out_features=3, bias=True)), ('linear2', Linear(in_features=3, out_features=4, bias=True)), ('batch_norm4', BatchNorm2d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))]