下載,保存,加載,使用模型權重
在這一節里面我們會過一遍對模型權重的常用操作,比如:
- 如何下載常用模型的預訓練權重
- 如何下載常用模型的無訓練權重(只下載網絡結構)
- 如何加載模型權重
- 如何保存權重
- 加載模型權重后進行推理的注意事項
寫在開頭:權重是以什么形式存在的?
光用,是肯定夠的,但是如果你能稍微懂一點原理,那么你很有可能在某一日突然融會貫通,做出非常牛逼的優化。
Pytorch模型會將學習到的參數存儲在稱為state_dict
的內部狀態字典中。為了深入探究,我們可以創建一個單線性層簡單模型,然后看看它的狀態字典長啥樣:
import torch
import torch.nn as nn
import jsonclass SimpleLinearModel(nn.Module):def __init__(self):super().__init__()self.linear = nn.Linear(5, 2) # 五個輸入,兩個輸出def forward(self, x):return self.linear(x)model = SimpleLinearModel()
model_state_dict = model.state_dict()
print(model_state_dict)
以上代碼塊的輸出如下:
雖然它是個字典,但是由于張量的存在讓他不是一個友好的鍵值對結構,而是元組結構,但是我們還是可以把他轉換JSON來直觀感受一下:
{"linear.weight": [[0.1941000074148178,0.22420001029968262,-0.3236999809741974,-0.1558000087738037,0.2337000072002411],[0.15130000114440918,0.11470000594854355,0.3953999876976013,-0.33970001339912415,-0.20650000870227814]],"linear.bias": [0.046799998730421066,-0.03530000150203705]
}
這下直觀多了,可以看到,state_dict
存儲了兩個學習的參數,其中包括了一個W2×5W_{2\times5}W2×5?的全連接矩陣和一個長度為2的偏置向量bbb。
不過需要注意的是這是一個有序字典,這樣才能保證數據在流經權重文件的時候才能一層一層的被處理。
下載并保存常用預訓練模型權重
torchvision.models
包內置了很多的不同任務模型權重,包括但不限于圖像分類,語義分割,實例分割,關鍵點檢測,視頻分類,光流等等,你可以逛逛這個權重菜市場,這里我就不放圖介紹了。
一般情況下,你可能要使用這些預訓練的模型來進行 遷移學習, pytorch加載這些權重相當簡單,一般的調用公式為:
model = model.模型名(weights='用什么數據集預訓練的')
舉個例子,我們加載一下vgg的在IMAGENET1K_V1
上的權重,看看它結構如何
model = models.vgg16(weights='IMAGENET1K_V1')
print(model)
可以看到這個包含預訓練權重的網絡模型已經被我們保存到state_dict
中了,但是它目前還只是在內存里面,沒有寫入外存(硬盤),如果你想把它保存到本地,你可以這樣做:
torch.save(model.state_dict(), 'vgg1-_model_weights.pth')
這樣就會把你當前的權重保存為一個.pth文件。
加載模型權重-僅加載權重
OK, 現在你已經有了一個vgg模型權重,那你要怎么把它加載到對應的網絡上呢?
正常的步驟是這樣的:
- 創建同一模型的實例 (不指定數據集的時候說明只要結構不要預訓練權重)
- 使用
load_state_dict()
方法加載參數
model = models.vgg16() # 這里沒有指定數據集,說明只要結構
model.load_state_dict(torch.load('model_weights.pth', weights_only=True)) # 這里加載權重,該.pth文件只包含權重,不包含結構。
注意,如果你使用
torch.save(model.state_dict(), 'path')
,只有權重會被保存!如果你想在保存權重的同時也保存模型結構,你可以這么做:torch.save(model, 'model.pth')
這個做法的優點是可以在加載被這樣保存的權重的時候無需初始化對應的網絡:
model = torch.load('model.pth', weights_only=False) # 說明這個model.pth并不是只保存了權重,還有模型架構,所以不需要先實例化再加載權重
但是,模型和權重一起加載并不是pytorch官方推薦的最佳實踐! pytorch官方推薦的方式還是只保存模型權重,要加載的時候先實例化網絡再加載權重(后一段就是講這個的)。這是因為.pth文件的解析基于pickle協議
實現,而pickle
文件不僅僅是數據存儲,它還可以包含可執行代碼。當 torch.load()
反序列化一個 pickle
文件時,它會執行文件中的字節碼來重新創建對象。
這代表:如果一個 .pth
文件是惡意創建的,它可能包含惡意代碼。當你在不知情的情況下加載這個文件時,這些惡意代碼會被執行,從而導致你的系統受到攻擊,例如被植入病毒、竊取數據等。weights_only=True
的作用就是切斷這個風險鏈。它告訴 PyTorch:“我只信任文件中的張量數據。不要執行任何其他的 Python 對象代碼,即使它們存在于文件中。”這樣就有效地防止了潛在的惡意代碼被執行。
使用預訓練權重進行推理
這里要說的不多,用預訓練權重加載好模型之后記得打開.eval
評估模式再開始推理:
model.eval()
.eval()
方法的作用是將模型切換到評估(evaluation)模式。這個模式會關閉一些在訓練時才需要的特殊層,以確保模型在推理時能夠產生一致且可預測的結果。
具體來說,他會關閉這兩個層的以下作用:
.eval()
是 PyTorch 模型推理時一個非常重要的步驟,你提到的這一點非常關鍵。
為什么需要在推理前調用 model.eval()
?
.eval()
方法的作用是將模型切換到評估(evaluation)模式。這個模式會關閉一些在訓練時才需要的特殊層,以確保模型在推理時能夠產生一致且可預測的結果。
具體來說,.eval()
主要影響以下兩種類型的層:
-
Dropout
層- 訓練模式 (
model.train()
):Dropout
層會以一定的概率隨機“丟棄”一些神經元的輸出,以防止模型過擬合。這意味著每次前向傳播(forward pass)時,網絡結構都是不一樣的。 - 評估模式 (
model.eval()
):Dropout
層會被關閉。所有神經元都參與計算,不再隨機丟棄。這確保了在推理時,每次對同一輸入進行預測,都會得到完全相同的結果。
- 訓練模式 (
-
BatchNorm
(批量歸一化)層- 訓練模式 (
model.train()
):BatchNorm
層會根據當前批次(batch)的輸入數據來計算均值和方差,并用這些統計量進行歸一化。 - 評估模式 (
model.eval()
):BatchNorm
層會停止更新均值和方差。它會使用在訓練階段已經學到的、全局的、固定的均值和方差來進行歸一化。這同樣是為了確保推理結果的穩定性,因為在推理時,我們通常只處理單個樣本或小批量的樣本,它們的統計量沒有代表性。
- 訓練模式 (
如果沒有調用 model.eval()
,模型將保持在訓練模式。這將導致:
- 結果不穩定:因為
Dropout
層會隨機丟棄神經元,即使輸入相同,每次推理的結果也可能不同。 - 結果不準確:
BatchNorm
層會使用不穩定的批次統計量進行歸一化,而不是使用訓練時學到的穩定統計量,這會導致推理結果的準確性下降。
所以,為了得到穩定、準確且可復現的推理結果,在使用預訓練模型進行預測時,必須在推理循環之前調用 model.eval()
。