Pytorch如何保存训练好的模型

0.为什么要保存和加载模型

用数据对模型进行训练后得到了比较理想的模型,但在实际应用的时候不可能每次都先进行训练然后再使用,所以就得先将之前训练好的模型保存下来,然后在需要用到的时候加载一下直接使用。模型的本质是一堆用某种结构存储起来的参数,所以在保存的时候有两种方式,一种方式是直接将整个模型保存下来,之后直接加载整个模型,但这样会比较耗内存;另一种是只保存模型的参数,之后用到的时候再创建一个同样结构的新模型,然后把所保存的参数导入新模型。

1.两种情况的实现方法

(1)只保存模型参数字典(推荐)

#保存
torch.save(the_model.state_dict(), PATH)
#读取
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

(2)保存整个模型

#保存
torch.save(the_model, PATH)
#读取
the_model = torch.load(PATH)

3.只保存模型参数的情况(例子)

pytorch会把模型的参数放在一个字典里面,而我们所要做的就是将这个字典保存,然后再调用。
比如说设计一个单层LSTM的网络,然后进行训练,训练完之后将模型的参数字典进行保存,保存为同文件夹下面的rnn.pt文件:

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers):
        super(LSTM, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, x):
        # Set initial states
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device) 
         # 2 for bidirection
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
        # Forward propagate LSTM
        out, _ = self.lstm(x, (h0, c0))  
        # out: tensor of shape (batch_size, seq_length, hidden_size*2)
        out = self.fc(out)
        return out


rnn = LSTM(input_size=1, hidden_size=10, num_layers=2).to(device)

# optimize all cnn parameters
optimizer = torch.optim.Adam(rnn.parameters(), lr=0.001)  
# the target label is not one-hotted
loss_func = nn.MSELoss()  

for epoch in range(1000):
    output = rnn(train_tensor)  # cnn output`
    loss = loss_func(output, train_labels_tensor)  # cross entropy loss
    optimizer.zero_grad()  # clear gradients for this training step
    loss.backward()  # backpropagation, compute gradients
    optimizer.step()  # apply gradients
    output_sum = output


# 保存模型
torch.save(rnn.state_dict(), 'rnn.pt')

保存完之后利用这个训练完的模型对数据进行处理:

# 测试所保存的模型
m_state_dict = torch.load('rnn.pt')
new_m = LSTM(input_size=1, hidden_size=10, num_layers=2).to(device)
new_m.load_state_dict(m_state_dict)
predict = new_m(test_tensor)

这里做一下说明,在保存模型的时候rnn.state_dict()表示rnn这个模型的参数字典,在测试所保存的模型时要先将这个参数字典加载一下m_state_dict = torch.load('rnn.pt')

然后再实例化一个LSTM对像,这里要保证传入的参数跟实例化rnn是传入的对象时一样的,即结构相同new_m = LSTM(input_size=1, hidden_size=10, num_layers=2).to(device)

下面是给这个新的模型传入之前加载的参数new_m.load_state_dict(m_state_dict)

最后就可以利用这个模型处理数据了predict = new_m(test_tensor)

4.保存整个模型的情况(例子)

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers):
        super(LSTM, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, x):
        # Set initial states
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)  # 2 for bidirection
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)

        # Forward propagate LSTM
        out, _ = self.lstm(x, (h0, c0))  # out: tensor of shape (batch_size, seq_length, hidden_size*2)
        # print("output_in=", out.shape)
        # print("fc_in_shape=", out[:, -1, :].shape)
        # Decode the hidden state of the last time step
        # out = torch.cat((out[:, 0, :], out[-1, :, :]), axis=0)
        # out = self.fc(out[:, -1, :])  # 取最后一列为out
        out = self.fc(out)
        return out


rnn = LSTM(input_size=1, hidden_size=10, num_layers=2).to(device)
print(rnn)


optimizer = torch.optim.Adam(rnn.parameters(), lr=0.001)  # optimize all cnn parameters
loss_func = nn.MSELoss()  # the target label is not one-hotted

for epoch in range(1000):
    output = rnn(train_tensor)  # cnn output`
    loss = loss_func(output, train_labels_tensor)  # cross entropy loss
    optimizer.zero_grad()  # clear gradients for this training step
    loss.backward()  # backpropagation, compute gradients
    optimizer.step()  # apply gradients
    output_sum = output


# 保存模型

torch.save(rnn, 'rnn1.pt')

保存完之后利用这个训练完的模型对数据进行处理:

new_m = torch.load('rnn1.pt')
predict = new_m(test_tensor)

参考pytorch的官方文档

  • 81
    点赞
  • 312
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 11
    评论
### 回答1: 要在 PyTorch保存训练好的模型,可以使用 PyTorch 的内置函数 `torch.save()`。下面是保存模型的基本步骤: 1. 在训练完成后,将模型保存为字典类型,包括模型的状态字典和其他必要的参数。 ```python model_state = { 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch, 'loss': loss } ``` 2. 使用 `torch.save()` 函数将模型保存到指定的文件路径。 ```python torch.save(model_state, 'model.pt') ``` 3. 若要加载模型,请使用 `torch.load()` 函数来加载状态字典并将其赋值给模型。 ```python model_state = torch.load('model.pt') model.load_state_dict(model_state['state_dict']) ``` 这些步骤可以适用于各种不同的 PyTorch 模型。另外,如果您想要更好地控制保存的内容和格式,可以使用 PyTorch 的序列化函数将模型保存为自定义格式的文件。 ### 回答2: 要保存训练好的模型,我们可以使用PyTorch的torch.save()函数。该函数接受两个参数:第一个是要保存模型的状态字典,第二个是保存路径。 首先,我们需要定义一个模型的实例,然后在训练过程中,使用optimizer进行模型的优化。当训练完成后,我们可以使用torch.save()函数将模型的状态字典保存下来。 下面是保存模型的示例代码: ```python import torch import torch.nn as nn import torch.optim as optim import torchvision.models as models # 定义一个模型实例 model = models.resnet18() # 定义优化器 optimizer = optim.Adam(model.parameters(), lr=0.001) # ... 在训练过程中进行模型训练 ... # 训练完成后,保存模型的状态字典 torch.save(model.state_dict(), 'model.pth') ``` 以上代码中,我们使用了ResNet18作为示例模型,Adam作为优化器。在训练完成后,使用torch.save()函数将模型的状态字典保存到'model.pth'文件中。 注意,只保存模型的状态字典是因为它包含了模型的参数和缓存信息,但不包含模型的结构。在加载模型时,我们需要先创建一个模型实例,然后使用torch.load()函数加载保存模型状态字典。接下来,通过调用模型实例的load_state_dict()方法加载模型的状态字典,使模型获取保存的参数与缓存信息。 希望以上回答对您有所帮助! ### 回答3: 在PyTorch保存训练好的模型可以通过以下步骤完成: 1. 导入必要的库: ``` import torch import torch.nn as nn import torch.optim as optim ``` 2. 定义模型结构: ``` class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() self.fc = nn.Linear(input_size, output_size) def forward(self, x): x = self.fc(x) return x ``` 3. 创建模型实例: ``` model = MyModel() ``` 4. 定义损失函数和优化器: ``` criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) ``` 5. 训练模型: ``` for epoch in range(num_epochs): # 前向传播 outputs = model(inputs) loss = criterion(outputs, labels) # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step() ``` 6. 保存模型: ``` torch.save(model.state_dict(), 'model.pth') ``` 此代码将模型的参数保存到名为`model.pth`的文件中。 7. 加载模型: ``` model = MyModel() model.load_state_dict(torch.load('model.pth')) model.eval() ``` 通过加载参数文件,我们可以将模型的状态恢复到训练完成后的状态,然后使用`model.eval()`将模型设置为评估模式。
评论 11
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

comli_cn

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值