Pytorch怎樣保存訓(xùn)練好的模型
為什么要保存和加載模型
用數(shù)據(jù)對模型進(jìn)行訓(xùn)練后得到了比較理想的模型,但在實際應(yīng)用的時候不可能每次都先進(jìn)行訓(xùn)練然后再使用,所以就得先將之前訓(xùn)練好的模型保存下來,然后在需要用到的時候加載一下直接使用。
模型的本質(zhì)是一堆用某種結(jié)構(gòu)存儲起來的參數(shù),所以在保存的時候有兩種方式
- 一種方式是直接將整個模型保存下來,之后直接加載整個模型,但這樣會比較耗內(nèi)存;
- 另一種是只保存模型的參數(shù),之后用到的時候再創(chuàng)建一個同樣結(jié)構(gòu)的新模型,然后把所保存的參數(shù)導(dǎo)入新模型。
兩種情況的實現(xiàn)方法
(1)只保存模型參數(shù)字典(推薦)
#保存 torch.save(the_model.state_dict(), PATH) #讀取 the_model = TheModelClass(*args, **kwargs) the_model.load_state_dict(torch.load(PATH))
(2)保存整個模型
#保存 torch.save(the_model, PATH) #讀取 the_model = torch.load(PATH)
只保存模型參數(shù)的情況(例子)
pytorch會把模型的參數(shù)放在一個字典里面,而我們所要做的就是將這個字典保存,然后再調(diào)用。
比如說設(shè)計一個單層LSTM的網(wǎng)絡(luò),然后進(jìn)行訓(xùn)練,訓(xùn)練完之后將模型的參數(shù)字典進(jìn)行保存,保存為同文件夾下面的rnn.pt文件:
class LSTM(nn.Module): ? ? def __init__(self, input_size, hidden_size, num_layers): ? ? ? ? super(LSTM, self).__init__() ? ? ? ? self.hidden_size = hidden_size ? ? ? ? self.num_layers = num_layers ? ? ? ? self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True) ? ? ? ? self.fc = nn.Linear(hidden_size, 1) ? ? def forward(self, x): ? ? ? ? # Set initial states ? ? ? ? h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)? ? ? ? ? ?# 2 for bidirection ? ? ? ? c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device) ? ? ? ? # Forward propagate LSTM ? ? ? ? out, _ = self.lstm(x, (h0, c0)) ? ? ? ? ? # out: tensor of shape (batch_size, seq_length, hidden_size*2) ? ? ? ? out = self.fc(out) ? ? ? ? return out rnn = LSTM(input_size=1, hidden_size=10, num_layers=2).to(device) # optimize all cnn parameters optimizer = torch.optim.Adam(rnn.parameters(), lr=0.001) ? # the target label is not one-hotted loss_func = nn.MSELoss() ? for epoch in range(1000): ? ? output = rnn(train_tensor) ?# cnn output` ? ? loss = loss_func(output, train_labels_tensor) ?# cross entropy loss ? ? optimizer.zero_grad() ?# clear gradients for this training step ? ? loss.backward() ?# backpropagation, compute gradients ? ? optimizer.step() ?# apply gradients ? ? output_sum = output # 保存模型 torch.save(rnn.state_dict(), 'rnn.pt')
保存完之后利用這個訓(xùn)練完的模型對數(shù)據(jù)進(jìn)行處理:
# 測試所保存的模型 m_state_dict = torch.load('rnn.pt') new_m = LSTM(input_size=1, hidden_size=10, num_layers=2).to(device) new_m.load_state_dict(m_state_dict) predict = new_m(test_tensor)
這里做一下說明,在保存模型的時候rnn.state_dict()表示rnn這個模型的參數(shù)字典,在測試所保存的模型時要先將這個參數(shù)字典加載一下
m_state_dict = torch.load('rnn.pt');
然后再實例化一個LSTM對像,這里要保證傳入的參數(shù)跟實例化rnn是傳入的對象時一樣的,即結(jié)構(gòu)相同
new_m = LSTM(input_size=1, hidden_size=10, num_layers=2).to(device);
下面是給這個新的模型傳入之前加載的參數(shù)
new_m.load_state_dict(m_state_dict);
最后就可以利用這個模型處理數(shù)據(jù)了
predict = new_m(test_tensor)
保存整個模型的情況(例子)
class LSTM(nn.Module): ? ? def __init__(self, input_size, hidden_size, num_layers): ? ? ? ? super(LSTM, self).__init__() ? ? ? ? self.hidden_size = hidden_size ? ? ? ? self.num_layers = num_layers ? ? ? ? self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True) ? ? ? ? self.fc = nn.Linear(hidden_size, 1) ? ? def forward(self, x): ? ? ? ? # Set initial states ? ? ? ? h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device) ?# 2 for bidirection ? ? ? ? c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device) ? ? ? ? # Forward propagate LSTM ? ? ? ? out, _ = self.lstm(x, (h0, c0)) ?# out: tensor of shape (batch_size, seq_length, hidden_size*2) ? ? ? ? # print("output_in=", out.shape) ? ? ? ? # print("fc_in_shape=", out[:, -1, :].shape) ? ? ? ? # Decode the hidden state of the last time step ? ? ? ? # out = torch.cat((out[:, 0, :], out[-1, :, :]), axis=0) ? ? ? ? # out = self.fc(out[:, -1, :]) ?# 取最后一列為out ? ? ? ? out = self.fc(out) ? ? ? ? return out rnn = LSTM(input_size=1, hidden_size=10, num_layers=2).to(device) print(rnn) optimizer = torch.optim.Adam(rnn.parameters(), lr=0.001) ?# optimize all cnn parameters loss_func = nn.MSELoss() ?# the target label is not one-hotted for epoch in range(1000): ? ? output = rnn(train_tensor) ?# cnn output` ? ? loss = loss_func(output, train_labels_tensor) ?# cross entropy loss ? ? optimizer.zero_grad() ?# clear gradients for this training step ? ? loss.backward() ?# backpropagation, compute gradients ? ? optimizer.step() ?# apply gradients ? ? output_sum = output # 保存模型 torch.save(rnn, 'rnn1.pt')
保存完之后利用這個訓(xùn)練完的模型對數(shù)據(jù)進(jìn)行處理:
new_m = torch.load('rnn1.pt') predict = new_m(test_tensor)
總結(jié)
以上為個人經(jīng)驗,希望能給大家一個參考,也希望大家多多支持腳本之家。
相關(guān)文章
Python cookbook(數(shù)據(jù)結(jié)構(gòu)與算法)實現(xiàn)對不原生支持比較操作的對象排序算法示例
這篇文章主要介紹了Python cookbook(數(shù)據(jù)結(jié)構(gòu)與算法)實現(xiàn)對不原生支持比較操作的對象排序算法,結(jié)合實例形式分析了Python針對類實例進(jìn)行排序相關(guān)操作技巧,需要的朋友可以參考下2018-03-03Python中列表和字符串常用的數(shù)據(jù)去重方法總結(jié)
關(guān)于數(shù)據(jù)去重,咱們這里簡單理解下,就是刪除掉重復(fù)的數(shù)據(jù),應(yīng)用的場景比如某些產(chǎn)品產(chǎn)生的大數(shù)據(jù),有很多重復(fù)的數(shù)據(jù),為了不影響分析結(jié)果,我們可能需要對這些數(shù)據(jù)進(jìn)行去重,所以本文給大家總結(jié)了Python中列表和字符串常用的數(shù)據(jù)去重方法,需要的朋友可以參考下2023-11-11利用python循環(huán)創(chuàng)建多個文件的方法
今天小編就為大家分享一篇利用python循環(huán)創(chuàng)建多個文件的方法,具有很好的參考價值,希望對大家有所幫助。一起跟隨小編過來看看吧2018-10-10Python+selenium點擊網(wǎng)頁上指定坐標(biāo)的實例
今天小編就為大家分享一篇Python+selenium點擊網(wǎng)頁上指定坐標(biāo)的實例,具有很好的參考價值,希望對大家有所幫助。一起跟隨小編過來看看吧2019-07-07python?操作?mongodb?數(shù)據(jù)庫詳情
這篇文章主要介紹了python?操作?mongodb?數(shù)據(jù)庫詳情,通過鏈接數(shù)據(jù)庫,創(chuàng)建數(shù)據(jù)庫展開內(nèi)容詳細(xì),具有一定的參考價值,需要的的小伙伴可以參考一下2022-04-04對dataframe數(shù)據(jù)之間求補集的實例詳解
今天小編就為大家分享一篇對dataframe數(shù)據(jù)之間求補集的實例詳解,具有很好的參考價值,希望對大家有所幫助。一起跟隨小編過來看看吧2019-01-01Python實戰(zhàn)購物車項目的實現(xiàn)參考
今天小編就為大家分享一篇關(guān)于Python實戰(zhàn)購物車項目的實現(xiàn)參考,小編覺得內(nèi)容挺不錯的,現(xiàn)在分享給大家,具有很好的參考價值,需要的朋友一起跟隨小編來看看吧2019-02-02解決Keras中循環(huán)使用K.ctc_decode內(nèi)存不釋放的問題
這篇文章主要介紹了解決Keras中循環(huán)使用K.ctc_decode內(nèi)存不釋放的問題,具有很好的參考價值,希望對大家有所幫助。一起跟隨小編過來看看吧2020-06-06