基於Pytorch的深度學習模型保存和加載方式_#開發語言


PyTorch2 Python深度學習 - 模型保存與加載

在PyTorch 2中，模型的保存與加載是一個常見的操作。通常，我們會使用torch.save()來保存模型的權重（即state_dict），並使用torch.load()來加載模型權重。以下是PyTorch 2中如何保存和加載模型的介紹及實例代碼：

1. 保存模型權重以及加載

保存模型的權重（state_dict）

# 保存模型權重
torch.save(model.state_dict(), 'model_weights.pt')

state_dict是PyTorch中保存模型的最佳方式，通常只保存模型的權重，而不包括模型結構。保存模型的權重後，你可以在加載時重新構建模型並加載權重。

from sklearn.datasets import load_iris
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import torch

# 加載鳶尾花數據集
iris = load_iris()
X = iris.data  # 150個樣本 4個特徵
y = iris.target  # 3個類別（Setosa, Versicolor, Virginica）

# 特徵標準化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 將數據集分為訓練集和測試集
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)

# 將數據轉換為Pytorch的Tensor格式
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.long)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.long)

# 創建數據集和數據加載器
train_dataset = torch.utils.data.TensorDataset(X_train, y_train)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataset = torch.utils.data.TensorDataset(X_test, y_test)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)

# 超參數定義
input_size = 4  # 輸入特徵數
hidden_size = 16  # 隱藏層節點數
output_size = 3  # 輸出類別數 （鳶尾花有3個類別）

# 創建模型
model = torch.nn.Sequential(
    torch.nn.Linear(input_size, hidden_size),  # 輸入層 -> 隱藏層
    torch.nn.ReLU(),
    torch.nn.Linear(hidden_size, output_size)  # 隱藏層 -> 輸出層（3類）
)

# 損失函數和優化器
criterion = torch.nn.CrossEntropyLoss()  # 交叉熵損失函數
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)  # 優化器

# 訓練模型
num_epochs = 100
for epoch in range(num_epochs):
    model.train()  # 設置為訓練模式
    for inputs, labels in train_loader:
        # 前向傳播
        outputs = model(inputs)
        loss = criterion(outputs, labels)  # 計算損失

        # 反向傳播和優化
        optimizer.zero_grad()  # 梯度清零
        loss.backward()  # 計算梯度
        optimizer.step()  # 更新參數

    # 輸出訓練結果
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 保存模型權重
torch.save(model.state_dict(), 'model_weights.pt')

在上面的代碼中，我們定義了一個簡單的全連接網絡，並保存了模型的權重到model_weights.pth文件。

加載模型權重（state_dict）

# 加載保存的模型權重
model.load_state_dict(torch.load('model_weights.pt'))

加載模型時，首先需要重新定義與保存時相同的模型結構，然後使用load_state_dict()將權重加載到模型中。

from sklearn.datasets import load_iris
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import torch

# 加載鳶尾花數據集
iris = load_iris()
X = iris.data  # 150個樣本 4個特徵
y = iris.target  # 3個類別（Setosa, Versicolor, Virginica）

# 特徵標準化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 將數據集分為訓練集和測試集
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)

# 將數據轉換為Pytorch的Tensor格式
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.long)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.long)

# 創建數據集和數據加載器
train_dataset = torch.utils.data.TensorDataset(X_train, y_train)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataset = torch.utils.data.TensorDataset(X_test, y_test)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)

# 超參數定義
input_size = 4  # 輸入特徵數
hidden_size = 16  # 隱藏層節點數
output_size = 3  # 輸出類別數 （鳶尾花有3個類別）

# 創建模型
model = torch.nn.Sequential(
    torch.nn.Linear(input_size, hidden_size),  # 輸入層 -> 隱藏層
    torch.nn.ReLU(),
    torch.nn.Linear(hidden_size, output_size)  # 隱藏層 -> 輸出層（3類）
)

# 加載保存的模型權重
model.load_state_dict(torch.load('model_weights.pt'))

# 測試模型
model.eval（)  # 設置為測試模式
y_pred = []
y_true = []
with torch.no_grad():  # 禁用梯度計算
    for inputs, labels in test_loader:
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)  # 獲取預測結果 獲取最大概率的索引
        y_pred.extend(predicted.numpy())
        y_true.extend(labels.numpy())

# 計算準確率
accuracy = accuracy_score(y_true, y_pred)
print(f'Accuracy: {accuracy * 100:.2f}%')

注意：model.eval（) 是 關鍵的一步，它會將 dropout 和 batch normalization 等層切換到推理模式，確保結果的一致性。

2. 保存整個模型（包括結構和權重）以及加載

除了保存模型的state_dict外，PyTorch還允許你保存整個模型（包括結構和權重）。不過這種方法不推薦，因為它依賴於特定的代碼實現，可能導致移植性問題。

# 保存模型權重
torch.save(model, 'whole_model.pt')

完整代碼：

from sklearn.datasets import load_iris
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import torch

# 加載鳶尾花數據集
iris = load_iris()
X = iris.data  # 150個樣本 4個特徵
y = iris.target  # 3個類別（Setosa, Versicolor, Virginica）

# 特徵標準化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 將數據集分為訓練集和測試集
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)

# 將數據轉換為Pytorch的Tensor格式
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.long)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.long)

# 創建數據集和數據加載器
train_dataset = torch.utils.data.TensorDataset(X_train, y_train)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataset = torch.utils.data.TensorDataset(X_test, y_test)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)

# 超參數定義
input_size = 4  # 輸入特徵數
hidden_size = 16  # 隱藏層節點數
output_size = 3  # 輸出類別數 （鳶尾花有3個類別）

# 創建模型
model = torch.nn.Sequential(
    torch.nn.Linear(input_size, hidden_size),  # 輸入層 -> 隱藏層
    torch.nn.ReLU(),
    torch.nn.Linear(hidden_size, output_size)  # 隱藏層 -> 輸出層（3類）
)

# 損失函數和優化器
criterion = torch.nn.CrossEntropyLoss()  # 交叉熵損失函數
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)  # 優化器

# 訓練模型
num_epochs = 100
for epoch in range(num_epochs):
    model.train()  # 設置為訓練模式
    for inputs, labels in train_loader:
        # 前向傳播
        outputs = model(inputs)
        loss = criterion(outputs, labels)  # 計算損失

        # 反向傳播和優化
        optimizer.zero_grad()  # 梯度清零
        loss.backward()  # 計算梯度
        optimizer.step()  # 更新參數

    # 輸出訓練結果
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 保存模型權重
torch.save(model, 'whole_model.pt')
# 加載保存的模型權重
model = torch.load('whole_model.pt', weights_only=False)

加載整個模型時，你可以直接使用torch.load()方法：

from sklearn.datasets import load_iris
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import torch

# 加載鳶尾花數據集
iris = load_iris()
X = iris.data  # 150個樣本 4個特徵
y = iris.target  # 3個類別（Setosa, Versicolor, Virginica）

# 特徵標準化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 將數據集分為訓練集和測試集
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)

# 將數據轉換為Pytorch的Tensor格式
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.long)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.long)

# 創建數據集和數據加載器
train_dataset = torch.utils.data.TensorDataset(X_train, y_train)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataset = torch.utils.data.TensorDataset(X_test, y_test)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)

# 加載保存的模型權重
model = torch.load('whole_model.pt', weights_only=False)

# 測試模型
y_pred = []
y_true = []
with torch.no_grad():  # 禁用梯度計算
    for inputs, labels in test_loader:
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)  # 獲取預測結果 獲取最大概率的索引
        y_pred.extend(predicted.numpy())
        y_true.extend(labels.numpy())

# 計算準確率
accuracy = accuracy_score(y_true, y_pred)
print(f'Accuracy: {accuracy * 100:.2f}%')