【PyTorch】(三)模型的创建、参数初始化、保存和加载
文章目录
- 1. 模型的创建
- 1.1. 创建方法
- 1.1.1. 通过使用模型组件
- 1.1.2. 通过继承nn.Module类
- 1.2. 模型组件
- 1.2.1. 网络层
- 1.2.2. 函数包
- 1.2.3. 容器
- 1.3. 将模型转移到GPU
- 2. 模型参数初始化
- 3. 模型的保存与加载
- 3.1. 只保存参数
- 3.2. 保存模型和参数
1. 模型的创建
1.1. 创建方法
1.1.1. 通过使用模型组件
可以直接使用模型组件快速创建模型。
import torch.nn as nnmodel = nn.Linear(10, 10),
print(model)
输出结果:
Linear(in_features=10, out_features=10, bias=True)
1.1.2. 通过继承nn.Module类
在__init__方法中使用模型组件定义模型各层。在forward方法中实现前向传播。
import torch.nn as nnclass Model(nn.Module):def __init__(self):super().__init__()self.layer1 = nn.Linear(10, 10)self.layer2 = nn.Linear(10, 10)self.layer3 = nn.Sequential(nn.Linear(10, 10),nn.ReLU(),nn.Linear(10, 10))def forward(self, x):x = self.layer1(x)x = self.layer2(x)x = self.layer3(x)return xmodel = Model()
print(model)
输出结果:
Model((layer1): Linear(in_features=10, out_features=10, bias=True)(layer2): Linear(in_features=10, out_features=10, bias=True)(layer3): Sequential((0): Linear(in_features=10, out_features=10, bias=True)(1): ReLU()(2): Linear(in_features=10, out_features=10, bias=True))
)
1.2. 模型组件
1.2.1. 网络层
1.2.2. 函数包
1.2.3. 容器
1.3. 将模型转移到GPU
方法与将数据转移到GPU类似,都有两种方法:
- model.to(device)
- mode.cuda()
import torch
import torch.nn as nn# 创建模型实例
model = nn.Sequential(nn.Linear(10, 10),nn.ReLU(),nn.Linear(10, 10)
)# 将模型移动到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# 也可以
model = model.cuda()
2. 模型参数初始化
3. 模型的保存与加载
模型保存和加载使用的python内置的pickle模块。
3.1. 只保存参数
import torch
import torch.nn as nn# 创建模型实例
model1 = nn.Sequential(nn.Linear(10, 10),nn.ReLU(),nn.Linear(10, 10)
)# 保存和加载参数
torch.save(model1.state_dict(), '../model/model_params.pkl')
model1.load_state_dict(torch.load('../model/model_params.pkl'))
3.2. 保存模型和参数
import torch
import torch.nn as nn# 创建模型实例
model1 = nn.Sequential(nn.Linear(10, 10),nn.ReLU(),nn.Linear(10, 10)
)# 保存和加载模型和参数
torch.save(model1, '../model/model.pt')
model2 = torch.load('../model/model.pt')
print(model2)