pytorch如何保存訓練模型參數(shù)并實現(xiàn)繼續(xù)訓練
最近的想法是在推薦模型中考慮根據(jù)用戶對推薦結(jié)果的后續(xù)選擇，利用已訓練的offline預訓練模型參數(shù)來更新新的結(jié)果。
簡單記錄一下中途保存參數(shù)和后續(xù)使用不同數(shù)據(jù)訓練的方法。
簡單模型和訓練數(shù)據(jù)
先準備一個簡單模型，簡單兩層linear出個分類結(jié)果。
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(64, 32)
self.linear1 = nn.Linear(32, 10)
self.relu = nn.ReLU()
def forward(self, x):
x = self.linear(x)
x = self.relu(x)
x = self.linear1(x)
return x準備訓練用數(shù)據(jù)，這里直接隨機兩份，同時給出配套的十個類的分類label結(jié)果。
要注意的是 crossEntropy 交叉熵只認 long 以上的tensor，label記得轉(zhuǎn)一下類型。
rand1 = torch.rand((100, 64)).to(torch.float)
label1 = np.random.randint(0, 10, size=100)
label1 = torch.from_numpy(label1).to(torch.long)
rand2 = torch.rand((100, 64)).to(torch.float)
label2 = np.random.randint(0, 10, size=100)
label2 = torch.from_numpy(label2).to(torch.long)訓練簡單使用交叉熵，優(yōu)化器Adam。
model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss = nn.CrossEntropyLoss()
iteration = 100
for i in range(iteration):
output = model(rand1)
my_loss = loss(output, label1)
optimizer.zero_grad()
my_loss.backward()
optimizer.step()
print("iteration:{} loss:{}".format(i, my_loss))反正能跑起來：
保存與讀取訓練參數(shù)結(jié)果的方法
關(guān)鍵的保存方法，可以分為兩種，一種是直接把模型整體保存：
torch.save(model, save_path)
兩個參數(shù)，模型和保存目錄。不過這種不常用，如果模型變化或者只需要其中一部分參數(shù)就不太靈活。
常用方法的是將需要的模型或優(yōu)化器參數(shù)取出以字典形式存儲，這樣可以在使用時初始化相關(guān)模型，讀入對應(yīng)參數(shù)即可。
def save_model(save_path, iteration, optimizer, model):
torch.save({'iteration': iteration,
'optimizer_dict': optimizer.state_dict(),
'model_dict': model.state_dict()},
save_path)
print("model save success")分別存儲訓練循環(huán)次數(shù)，優(yōu)化器設(shè)置和模型參數(shù)結(jié)果。
初始化模型，讀取參數(shù)并設(shè)置：
def load_model(save_name, optimizer, model):
model_data = torch.load(save_name)
model.load_state_dict(model_data['model_dict'])
optimizer.load_state_dict(model_data['optimizer_dict'])
print("model load success")初始化新模型：
path = "net.dict"
save_model(path, iteration, optimizer, model)
print(model.state_dict()['linear.weight'])
new_model = MyModel()
new_optimizer = torch.optim.Adam(new_model.parameters(), lr=0.01)
load_model(path, new_optimizer, new_model)
print(new_model.state_dict()['linear.weight'])輸出第一個linear層的參數(shù)看看，確實相同，參數(shù)成功讀取上了。注意optimizer的初始化對應(yīng)模型別寫錯了。
之后用新模型繼續(xù)訓練試試：
for i in range(iteration):
output = new_model(rand2)
my_loss = loss(output, label2)
new_optimizer.zero_grad()
my_loss.backward()
new_optimizer.step()
print("iteration:{} loss:{}".format(i, my_loss))能成功訓練。
變化學習率的保存
上面的demo只用了固定的學習率來做實驗。
如果使用了 scheduler 來變化步長，只要保存 scheduler 的 state_dict ，之后對新初始化的 scheduler 設(shè)置對應(yīng)的當前循環(huán)次數(shù)即可。
# 存儲時 'scheduler': scheduler.state_dict() # 讀取時 scheduler.load_state_dict(checkpoint['lr_schedule'])
scheduler的使用可以看看我之前整理的文章：利用scheduler實現(xiàn)learning-rate學習率動態(tài)變化
總結(jié)
這次主要是整理了一下pytorch模型參數(shù)的整體保存方法，來實現(xiàn)新數(shù)據(jù)的后續(xù)訓練或直接作為offline預訓練模型來使用。
不過后續(xù)數(shù)據(jù)分布不同的話感覺效果會很差啊…
也不知道能不能用什么算法修改下權(quán)重來貼合新的數(shù)據(jù)，找點多次訓練優(yōu)化論文看看好了。
以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持腳本之家。
相關(guān)文章
這篇文章主要介紹了Python讀寫文件模式和文件對象方法,結(jié)合實例形式詳細分析了Python文件操作常用技巧與相關(guān)注意事項,需要的朋友可以參考下2019-09-09
這篇文章主要介紹了python 如何將IP 地址轉(zhuǎn)整數(shù)，幫助大家了解轉(zhuǎn)換的原理與收益，更好的理解python，感興趣的朋友可以了解下2020-11-11
這篇文章主要為大家詳細介紹了如何使用Python Tkinter編寫一個學生姓名比較工具,文中的示例代碼講解詳細,感興趣的小伙伴可以跟隨小編一起學習一下2025-10-10
python中np.random.permutation函數(shù)實例詳解
np.random.permutation是numpy中的一個函數(shù),它可以將一個數(shù)組中的元素隨機打亂,返回一個打亂后的新數(shù)組,下面這篇文章主要給大家介紹了關(guān)于python中np.random.permutation函數(shù)的相關(guān)資料,需要的朋友可以參考下2023-04-04
這篇文章主要為大家詳細介紹了如何利用Python語言制作簡易的核酸檢測日歷，文中的示例代碼講解詳細，感興趣的小伙伴可以動手嘗試一下2022-09-09
Python自定義一個類實現(xiàn)字典dict功能的方法
今天小編就為大家分享一篇Python自定義一個類實現(xiàn)字典dict功能的方法，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧2019-01-01
這篇文章主要介紹了Django中使用Celery的方法示例，小編覺得挺不錯的，現(xiàn)在分享給大家，也給大家做個參考。一起跟隨小編過來看看吧2018-11-11
Python常見讀寫文件操作實例總結(jié)【文本、json、csv、pdf等】
這篇文章主要介紹了Python常見讀寫文件操作,結(jié)合實例形式總結(jié)分析了Python常見的各種文件讀寫操作,包括文本、json、csv、pdf等文件的讀寫與相關(guān)注意事項,需要的朋友可以參考下2019-04-04
這篇文章主要介紹了詳解Django的CSRF認證實現(xiàn)，詳細的介紹了csrf原理和實現(xiàn)，小編覺得挺不錯的，現(xiàn)在分享給大家，也給大家做個參考。一起跟隨小編過來看看吧2018-10-10
在日常數(shù)據(jù)處理的過程中，有時候拿到手的是完整的地址信息，如果需要從地址信息中解析出相應(yīng)的各級行政單元名稱，方式有很多，而今天要介紹的方式只需要一行代碼即可快速實現(xiàn)，快跟隨小編一起學習一下吧2022-05-05