0705-深度網路模型持久化
pytorch完整教程目錄:https://www.cnblogs.com/nickchen121/p/14662511.html
一、持久化概述
在 torch 中,以下物件可以持久化到硬碟,並可以通過相應的方法把這些物件持久化到記憶體中:
- Tensor
- Variable
- nn.Module
- Optimizer
上述物件本質上最後都是儲存為 Tensor。並且 Tensor 的儲存和載入非常簡單,使用 t.save
和 t.load
即可。
在 save/load 時可指定使用的 pickle 模組,在 load 時還可以把 GPU tensor 對映到 CPU 或者其他 GPU 上。
我們可以通過 t.save(obj, file_name)
儲存任意可序列化的物件,然後通過 obj=t.load(file_name)
方法載入儲存的資料。
對於 Module 和 Optimizer 物件,建議儲存為對應的 state_dict,而不是直接儲存整個 Module/Optimizer 物件。Optimizer 物件儲存的是引數和動量資訊,通過載入之前的動量資訊,能夠很有效地減少模型震盪。
二、tensor 物件的儲存和載入
import torch as t
a = t.Tensor(3, 4)
if t.cuda.is_available():
a = a.cuda(1) # 把 a 轉為 GPU1 上的 tensor
t.save(a, 'a.pth')
# 載入為 b,儲存於 GPU1 上(因為儲存時 tensor 就在 GPU1 上)
b = t.load('a.pth')
# 載入為 c,儲存於 CPU
c = t.load('a.pth', map_location=lambda storage, loc: storage)
# 載入為 d,儲存於 GPU0 上
d = t.load('a.pth', map_location={'cuda:1': 'cuda:0'})
三、Module 物件的儲存和載入
t.set_default_tensor_type('torch.FloatTensor')
from torchvision.models import AlexNet
model = AlexNet()
# module 的 state_dict 是一個字典
model.state_dict().keys()
t.save(model.state_dict(), 'alexnet.pth')
model.load_state_dict(t.load('alexnet.pth'))
<All keys matched successfully>
四、Optimizer 物件的儲存和載入
optimizer = t.optim.Adam(model.parameters(), lr=0.1)
t.save(optimizer.state_dict(), 'optimizer.pth')
optimizer.load_state_dict(t.load('optimizer.pth'))
五、所有物件集合的儲存和載入
all_data = dict(optimizer=optimizer.state_dict(),
model=model.state_dict(),
info=u'模型和優化器的所有引數')
t.save(all_data, 'all.pth')
all_data = t.load('all.pth')
all_data.keys()
dict_keys(['optimizer', 'model', 'info'])
六、第七章總結
本章介紹了 torch 的很多工具模組,主要涉及資料載入、視覺化和 GPU 加速相關的內容,合理地使用這些模組可以極大地提升我們的編碼效率。