PyTorch學習：載入模型和引數
PyTorch學習：載入模型和引數
目錄pytorch的模型和引數是分開的，可以分別儲存或載入模型和引數。
pytorch有兩種模型儲存方式：
-
儲存整個神經網路的的結構資訊和模型引數資訊，save的物件是網路net
-
只儲存神經網路的訓練模型引數，save的物件是net.state_dict()
對應兩種儲存模型的方式，pytorch也有兩種載入模型的方式：
-
第一種儲存方式，載入模型時通過torch.load('.pth')直接初始化新的神經網路物件；
-
第二種儲存方式，需要首先匯入對應的網路，再通過net.load_state_dict(torch.load('.pth'))完成模型引數的載入。
在網路比較大的時候，第一種方法會花費較多的時間。
1. 直接載入模型和引數
載入別人訓練好的模型：
# 儲存和載入整個模型
torch.save(net, 'net.pth')
model = torch.load('net.pth')
2. 分別載入網路的結構和引數
# 將my_resnet模型儲存為my_resnet.pth torch.save(net.state_dict(), "net_parameter.pth") # 載入resnet，模型存放在my_resnet.pth net.load_state_dict(torch.load("net_parameter.pth"))
其中my_resnet是my_resnet.pth對應的網路結構。
3. pytorch預訓練模型
1）載入預訓練模型和引數
resnet18 = models.resnet18(pretrained=True)
2）只加載模型，不載入預訓練引數
# 匯入模型結構
resnet18 = models.resnet18(pretrained=False)
# 載入預先下載好的預訓練引數到resnet18
resnet18.load_state_dict(torch.load('resnet18-5c106cde.pth'))
3）載入部分預訓練模型
resnet152 = models.resnet152(pretrained=True) pretrained_dict = resnet152.state_dict() """載入torchvision中的預訓練模型和引數後通過state_dict()方法提取引數 也可以直接從官方model_zoo下載： pretrained_dict = model_zoo.load_url(model_urls['resnet152'])""" model_dict = model.state_dict() # 將pretrained_dict裡不屬於model_dict的鍵剔除掉 pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} # 更新現有的model_dict model_dict.update(pretrained_dict) # 載入我們真正需要的state_dict model.load_state_dict(model_dict)
這裡是直接呼叫pytorch中的常用模型
# PyTorch中的torchvision裡有很多常用的模型，可以直接呼叫：
import torchvision.models as models
resnet101 = models.resnet18()
alexnet = models.alexnet()
squeezenet = models.squeezenet1_0()
densenet = models.densenet_161()
序列化 Serialization
torch.saves[source]
torch.save(obj, f, pickle_module=<module 'pickle' from '/home/jenkins/miniconda/lib/python3.5/pickle.py'>, pickle_protocol=2)
儲存一個物件到一個硬碟檔案上 參考: Recommended approach for saving a model 引數：
- obj – 儲存物件
- f － 類檔案物件 (返回檔案描述符)或一個儲存檔名的字串
- pickle_module – 用於pickling元資料和物件的模組
- pickle_protocol – 指定pickle protocal 可以覆蓋預設引數
torch.load[source]
torch.load(f, map_location=None, pickle_module=<module 'pickle' from '/home/jenkins/miniconda/lib/python3.5/pickle.py'>)
從磁碟檔案中讀取一個通過torch.save()儲存的物件。 torch.load() 可通過引數map_location 動態地進行記憶體重對映，使其能從不動裝置中讀取檔案。一般呼叫時，需兩個引數: storage 和 location tag. 返回不同地址中的storage，或著返回None (此時地址可以通過預設方法進行解析). 如果這個引數是字典的話，意味著其是從檔案的地址標記到當前系統的地址標記的對映。 預設情況下， location tags中 "cpu"對應host tensors，‘cuda:device_id’ (e.g. ‘cuda:2’) 對應cuda tensors。 使用者可以通過register_package進行擴充套件，使用自己定義的標記和反序列化方法。
引數:
- f – 類檔案物件 (返回檔案描述符)或一個儲存檔名的字串
- map_location – 一個函式或字典規定如何remap儲存位置
- pickle_module – 用於unpickling元資料和物件的模組 (必須匹配序列化檔案時的pickle_module )
例子:
>>> torch.load('tensors.pt')
# Load all tensors onto the CPU
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# Map tensors from GPU 1 to GPU 0
>>> torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})
其他
關於.pt、.pth與.pkl 的區別
不存在格式上的區別，只是字尾名不同而已。在用torch.save()函式儲存模型檔案的時候，有些人喜歡用.pt字尾，有些人喜歡用.pth或 .pkl，用相同的 torch.save()語句儲存出來的模型檔案沒有什麼不同。
為什麼會有 .pkl這種字尾名呢？因為Python有一個序列化模組 - pickle，使用它儲存模型時，通常會起一個以 .pkl為字尾名的檔案。torch.save()正是使用pickle來儲存模型的。
關於Python序列化
我們把物件（變數）從記憶體中轉換為可儲存或傳輸的格式的過程稱之為序列化，在python中稱為pickle。
序列化之後，就可以把序列化後的內容寫入磁碟，或者通過網路傳輸到別的機器上（因為硬碟或網路傳輸時只接受bytes）。
反過來，把變數內容從序列化的物件重新讀到記憶體裡稱之為反序列化，即unpacking。
在python中，可以使用pickle和json兩個模組對資料進行序列化操作
其中：
- json可以用於字串或者字典等與python資料型別之間的序列化與反序列化操作
- pickle可以用於python特有型別與python資料型別之間的序列化與反序列化操作
參考資料：
PyTorch學習：載入模型和引數_lscelory的部落格-CSDN部落格_pytorch載入模型引數
torch - PyTorch中文文件 (pytorch-cn.readthedocs.io)
[原創] PyTorch模型 .pt，.pth，.pkl 的區別 – 編碼無悔 / Intent & Focused (codelast.com)