pytorch怎么保存最佳模型

小億
105
2024-06-04 17:10:28

在PyTorch中保存最佳模型通常是通過保存模型的參數(shù)和優(yōu)化器狀態(tài)來實(shí)現(xiàn)的。以下是一個(gè)示例代碼，演示了如何保存最佳模型：

import torch
import torch.nn as nn
import torch.optim as optim

# 定義模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

model = MyModel()

# 定義損失函數(shù)和優(yōu)化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 訓(xùn)練模型
best_loss = float('inf')
for epoch in range(num_epochs):
    # 訓(xùn)練過程
    train_loss = 0.0
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()

    train_loss /= len(train_loader)

    # 保存最佳模型
    if train_loss < best_loss:
        best_loss = train_loss
        torch.save({
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'best_loss': best_loss
        }, 'best_model.pth')

在上面的示例中，我們首先定義了一個(gè)模型、損失函數(shù)和優(yōu)化器。然后在訓(xùn)練過程中，我們通過比較當(dāng)前訓(xùn)練損失和最佳損失來保存最佳模型。當(dāng)訓(xùn)練損失小于最佳損失時(shí)，我們保存模型的狀態(tài)字典和優(yōu)化器的狀態(tài)字典，并將最佳損失更新為當(dāng)前訓(xùn)練損失。

最后，我們可以通過加載best_model.pth文件來恢復(fù)最佳模型的狀態(tài)，并繼續(xù)使用該模型進(jìn)行推理或進(jìn)一步的訓(xùn)練。

0 踩