科技猛獸：PyTorch 50.Pytorch模型保存與加載，并在加載的模型基礎(chǔ)上繼續(xù)訓(xùn)練?zhuanlan.zhihu.com

1.Pytorch 模型保存與加載，并在加載的模型基礎(chǔ)上繼續(xù)訓(xùn)練

只保存參數(shù)；（官方推薦）

torch.save(model.state_dict(), path) #只保存參數(shù) #特別地，如果還想保存某一次訓(xùn)練采用的優(yōu)化器、epochs等信息，可將這些信息組合 #起來構(gòu)成一個字典，然后將字典保存起來 state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch} torch.save(state, path)#加載 model.load_state_dict(torch.load(path))#針對上述第二種以字典形式保存的方法，加載方式如下 checkpoint = torch.load(path) model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) epoch = checkpoint(['epoch'])

ps:只保存參數(shù)的方法在加載的時候要事先定義好跟原模型一致的模型，并在該模型的實例對象(假設(shè)名為model)上進(jìn)行加載，即在使用上述加載語句前已經(jīng)有定義了一個和原模型一樣的Net, 并且進(jìn)行了實例化 model=Net( )

下面給出一個具體的例子程序，該程序只保存最新的參數(shù)

#-*- coding:utf-8 -*-import torch as torch import torchvision as tv import torch.nn as nn import torch.optim as optim import torch.nn.functional as F import torchvision.transforms as transforms from torchvision.transforms import ToPILImage import torch.backends.cudnn as cudnn import datetime import argparse# 參數(shù)聲明 batch_size = 32 epochs = 10 WORKERS = 0 # dataloder線程數(shù) test_flag = True #測試標(biāo)志，True時加載保存好的模型進(jìn)行測試 ROOT = '/home/pxt/pytorch/cifar' log_dir = '/home/pxt/pytorch/logs/cifar_model.pth' # 模型保存路徑# 加載MNIST數(shù)據(jù)集 transform = tv.transforms.Compose([transforms.ToTensor(),transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])train_data = tv.datasets.CIFAR10(root=ROOT, train=True, download=True, transform=transform) test_data = tv.datasets.CIFAR10(root=ROOT, train=False, download=False, transform=transform)train_load = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=WORKERS) test_load = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=WORKERS)# 構(gòu)造模型 class Net(nn.Module):def __init__(self):super(Net, self).__init__()self.conv1 = nn.Conv2d(3, 64, 3, padding=1)self.conv2 = nn.Conv2d(64, 128, 3, padding=1)self.conv3 = nn.Conv2d(128, 256, 3, padding=1)self.conv4 = nn.Conv2d(256, 256, 3, padding=1)self.pool = nn.MaxPool2d(2, 2)self.fc1 = nn.Linear(256 * 8 * 8, 1024)self.fc2 = nn.Linear(1024, 256)self.fc3 = nn.Linear(256, 10)def forward(self, x):x = F.relu(self.conv1(x))x = self.pool(F.relu(self.conv2(x)))x = F.relu(self.conv3(x))x = self.pool(F.relu(self.conv4(x)))x = x.view(-1, x.size()[1] * x.size()[2] * x.size()[3])x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return xdevice = torch.device("cuda" if use_cuda else "cpu") #model = Net().cuda() model = Net().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.01)# 模型訓(xùn)練 def train(model, train_loader, epoch):model.train()train_loss = 0for i, data in enumerate(train_loader, 0):x, y = datax = x.to(device)y = y.to(device)optimizer.zero_grad()y_hat = model(x)loss = criterion(y_hat, y)loss.backward()optimizer.step()train_loss += lossloss_mean = train_loss / (i+1)print('Train Epoch: {}t Loss: {:.6f}'.format(epoch, loss_mean.item()))# 模型測試 def test(model, test_loader):model.eval()test_loss = 0correct = 0with torch.no_grad():for i, data in enumerate(test_loader, 0):x, y = datax = x.cuda()y = y.cuda()optimizer.zero_grad()y_hat = model(x)test_loss += criterion(y_hat, y).item()pred = y_hat.max(1, keepdim=True)[1]correct += pred.eq(y.view_as(pred)).sum().item()test_loss /= (i+1)print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)n'.format(test_loss, correct, len(test_data), 100. * correct / len(test_data)))def main():# 如果test_flag=True,則加載已保存的模型if test_flag:# 加載保存的模型直接進(jìn)行測試機驗證，不進(jìn)行此模塊以后的步驟checkpoint = torch.load(log_dir)model.load_state_dict(checkpoint['model'])optimizer.load_state_dict(checkpoint['optimizer'])epochs = checkpoint['epoch']test(model, test_load)returnfor epoch in range(0, epochs):train(model, train_load, epoch)test(model, test_load)# 保存模型state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}torch.save(state, log_dir)if __name__ == '__main__':main()

在加載的模型基礎(chǔ)上繼續(xù)訓(xùn)練

在訓(xùn)練模型的時候可能會因為一些問題導(dǎo)致程序中斷，或者常常需要觀察訓(xùn)練情況的變化來更改學(xué)習(xí)率等參數(shù)，這時候就需要加載中斷前保存的模型，并在此基礎(chǔ)上繼續(xù)訓(xùn)練，這時候只需要對上例中的 main() 函數(shù)做相應(yīng)的修改即可，修改后的 main() 函數(shù)如下：

def main():# 如果test_flag=True,則加載已保存的模型if test_flag:# 加載保存的模型直接進(jìn)行測試機驗證，不進(jìn)行此模塊以后的步驟checkpoint = torch.load(log_dir)model.load_state_dict(checkpoint['model'])optimizer.load_state_dict(checkpoint['optimizer'])start_epoch = checkpoint['epoch']test(model, test_load)return# 如果有保存的模型，則加載模型，并在其基礎(chǔ)上繼續(xù)訓(xùn)練if os.path.exists(log_dir):checkpoint = torch.load(log_dir)model.load_state_dict(checkpoint['model'])optimizer.load_state_dict(checkpoint['optimizer'])start_epoch = checkpoint['epoch']print('加載 epoch {} 成功！'.format(start_epoch))else:start_epoch = 0print('無保存模型，將從頭開始訓(xùn)練！')for epoch in range(start_epoch+1, epochs):train(model, train_load, epoch)test(model, test_load)# 保存模型state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}torch.save(state, log_dir)

2. tqdm 進(jìn)度條 pip install tpdm

import time from tqdm import *for i in tqdm(range(1000)): # 傳入可迭代對象time.sleep(.01)pbar = tqdm(["a", "b", "c", "d"]) for char in pbar: # 可以為進(jìn)度條設(shè)置描述# 設(shè)置描述pbar.set_description("Processing %s" % char)time.sleep(1)

3. lmdb pip install lmdb

趙劍行：lmdb 數(shù)據(jù)庫?zhuanlan.zhihu.com#LMDB的全稱是Lightning Memory-Mapped Database(快如閃電的內(nèi)存映射數(shù)據(jù)庫) import lmdb import os, sysdef initialize():env = lmdb.open("lmdb_dir")return envdef insert(env, sid, name):txn = env.begin(write=True)txn.put(str(sid).encode(), name.encode())txn.commit()def delete(env, sid):txn = env.begin(write=True)txn.delete(str(sid).encode())txn.commit()def update(env, sid, name):txn = env.begin(write=True)txn.put(str(sid).encode(), name.encode())txn.commit()def search(env, sid):txn = env.begin()name = txn.get(str(sid).encode()).decode()return namedef display(env):txn = env.begin()cur = txn.cursor()for key, value in cur:print(key.decode(), value.decode())env = initialize()print("Insert 3 records.") insert(env, 1, "Alice") insert(env, 2, "Bob") insert(env, 3, "Peter") display(env)print("Delete the record where sid = 1.") delete(env, 1) display(env)print("Update the record where sid = 3.") update(env, 3, "Mark") display(env)print("Get the name of student whose sid = 3.") name = search(env, 3) print(name)#------------------------------------ """ Insert 3 records. 1 Alice 2 Bob 3 Peter Delete the record where sid = 1. 2 Bob 3 Peter Update the record where sid = 3. 2 Bob 3 Mark Get the name of student whose sid = 3. Mark """

3. PyYaml pip install pyyaml

Pytorch日積月累?www.zhihu.com# -*- coding:utf-8 -*- import yaml#https://www.zhihu.com/column/c_1244271799468482560 f = open(r'dataconfig.yml') dic = yaml.load(f, Loader=yaml.FullLoader) print (type(dic), dic) # type(dic) <class 'dict'>dic2 = {'name': 'Tom Smith', 'age': 37, 'spouse': {'name': 'Jane Smith', 'age': 25},'children': [{'name': 'Jimmy Smith', 'age': 15},{'name1': 'Jenny Smith', 'age1': 12}]} f_s = open(r'dataconfig2.yml','w') yaml.dump(dic2, f_s)

4. Pytorch中的Conv1d()和Conv2d()函數(shù)

Pytorch中的Conv1d()和Conv2d()函數(shù)_潘多拉星系的專欄-CSDN博客?blog.csdn.netpytorch之nn.Conv1d詳解_sunny_xsc1994的專欄-CSDN博客?blog.csdn.netPyTorch中的nn.Conv1d與nn.Conv2d?www.jianshu.com

一般來說，一維卷積nn.Conv1d()用于文本數(shù)據(jù)，只對寬度進(jìn)行卷積，對高度不卷積。通常，輸入大小為word_embedding_dim * max_sent_length，其中，word_embedding_dim為詞向量的維度，max_sent_length為句子的最大長度。卷積核窗口在句子長度的方向上滑動，進(jìn)行卷積操作。

輸入：批大小為50，句子的最大長度為35，詞向量維度為300

目標(biāo)：句子分類，共2類

import torch import torch.nn as nn# max_sent_len=35, batch_size=50, embedding_size=300 conv1 = nn.Conv1d(in_channels=300, out_channels=100, kernel_size=3) input = torch.randn(50, 35, 300) # batch_size x max_sent_len x embedding_size -> batch_size x embedding_size x max_sent_len input = input.permute(0, 2, 1) print("input:", input.size()) output = conv1(input) print("output:", output.size()) # 最大池化 pool1d = nn.MaxPool1d(kernel_size=35-3+1) pool1d_value = pool1d(output) print("最大池化輸出：", pool1d_value.size()) # 全連接 fc = nn.Linear(in_features=100, out_features=2) fc_inp = pool1d_value.view(-1, pool1d_value.size(1)) print("全連接輸入：", fc_inp.size()) fc_outp = fc(fc_inp) print("全連接輸出：", fc_outp.size()) # softmax m = nn.Softmax() out = m(fc_outp) print("輸出結(jié)果值：", out)

• 原始輸入大小為(50, 35, 300)，經(jīng)過permute(0, 2, 1)操作后，輸入的大小變?yōu)?50, 300, 35)；
• 使用1個window_size為3的卷積核進(jìn)行卷積，因為一維卷積是在最后維度上掃的，最后output的大小即為：50*100*（35-3+1）=50*100*33
• output經(jīng)過最大池化操作后，得到了數(shù)據(jù)維度為：(50,100,1)
• 經(jīng)過（輸入特征=100，輸出特征=2）的全連接層，數(shù)據(jù)維度就變?yōu)榱?#xff1a;(50，2)
• 再經(jīng)過softmax函數(shù)就得到了屬于兩個類別的概率值

eg: 假設(shè) embedding_size=256, feature_size=100, window_sizes=[3,4,5,6], max_text_len=35

class TextCNN(nn.Module):def __init__(self, config):super(TextCNN, self).__init__()self.is_training = Trueself.dropout_rate = config.dropout_rateself.num_class = config.num_classself.use_element = config.use_elementself.config = configself.embedding = nn.Embedding(num_embeddings=config.vocab_size, embedding_dim=config.embedding_size)self.convs = nn.ModuleList([nn.Sequential(nn.Conv1d(in_channels=config.embedding_size, out_channels=config.feature_size, kernel_size=h), # nn.BatchNorm1d(num_features=config.feature_size), nn.ReLU(),nn.MaxPool1d(kernel_size=config.max_text_len-h+1))for h in config.window_sizes ])self.fc = nn.Linear(in_features=config.feature_size*len(config.window_sizes),out_features=config.num_class)if os.path.exists(config.embedding_path) and config.is_training and config.is_pretrain:print("Loading pretrain embedding...")self.embedding.weight.data.copy_(torch.from_numpy(np.load(config.embedding_path))) def forward(self, x):embed_x = self.embedding(x)#print('embed size 1',embed_x.size()) # 32*35*256 # batch_size x text_len x embedding_size -> batch_size x embedding_size x text_lenembed_x = embed_x.permute(0, 2, 1)#print('embed size 2',embed_x.size()) # 32*256*35out = [conv(embed_x) for conv in self.convs] #out[i]:batch_size x feature_size*1#for o in out:# print('o',o.size()) # 32*100*1out = torch.cat(out, dim=1) # 對應(yīng)第二個維度（行）拼接起來，比如說5*2*1,5*3*1的拼接變成5*5*1#print(out.size(1)) # 32*400*1out = out.view(-1, out.size(1)) #print(out.size()) # 32*400 if not self.use_element:out = F.dropout(input=out, p=self.dropout_rate)out = self.fc(out)return out

基于上述代碼，具體計算過程如下：

原始輸入大小為(32, 35, 256)，經(jīng)過permute(0, 2, 1)操作后，輸入的大小變?yōu)?32, 256, 35)；

使用1個卷積核進(jìn)行卷積，可得到1個大小為32 x 100 x 1的輸出，共4個卷積核，故共有4個大小為32 x 100 x 1的輸出；

將上一步得到的4個結(jié)果在dim = 1上進(jìn)行拼接，輸出大小為32 x 400 x 1；

view操作后，輸出大小變?yōu)?2 x 400；

全連接，最終輸出大小為32 x 2，即分別預(yù)測為2類的概率大小。

鳶尾花分類 conv1 簡單demo

import numpy as np import torch import torch.nn as nn from sklearn.datasets import load_iris import torch.nn.functional as Fdata = load_iris()X = data.data # 只包括樣本的特征，150x4 y = data.target # 樣本的類型，[0, 1, 2] # print(X.shape) # (150, 4) # print(X[0]) #[5.1 3.5 1.4 0.2]input = X.reshape(150, 4, 1) # batch_size embedding_size max_sent_len input = torch.tensor(input, dtype=torch.float32) target = torch.LongTensor(y)# conv1 class conv1_model(nn.Module):"""conv1 文本分類建模 """def __init__(self):super(conv1_model, self).__init__()self.conv1 = nn.Conv1d(in_channels=4, out_channels=10, kernel_size=1)self.fc = nn.Linear(in_features=10, out_features=3)def forward(self, x):x = self.conv1(x)x = x.view(-1, x.size(1))out = self.fc(x)return outmodel = conv1_model()cost = nn.CrossEntropyLoss() #損失函數(shù)# 優(yōu)化器Adam優(yōu)化算法，相比SGD可以自動調(diào)節(jié)學(xué)習(xí)率 optimizer = torch.optim.Adam(model.parameters(), lr=0.01)for epoch in range(300):#x, y = [input, target] # 讀入一批的數(shù)據(jù)for i in range(1, 16):x, y = input[(i-1)*10:i*10], target[(i-1)*10:i*10]optimizer.zero_grad()# 模型預(yù)測predict = model(x)#print(predict.size(), target.size())loss = cost(predict, y)#print("loss ----> ", loss.item())loss.backward()optimizer.step()if epoch % 50 ==0:# [4.8 3.1 1.6 0.2] [5. 2. 3.5 1. ] [6.9 3.2 5.7 2.3] --> [0, 1, 2]pre_x = np.array([[4.8, 3.1, 1.6, 0.2], [5., 2., 3.5, 1. ], [6.9, 3.2, 5.7, 2.3]]).reshape(3, 4, 1)input_x = torch.tensor(pre_x, dtype=torch.float32)res = model(input_x)c = F.softmax(res, dim=1).argmax(dim=1)print("epoch, res ---> ",epoch, c)

總結(jié)

以上是生活随笔為你收集整理的pytorch保存模型pth_Pytorch_trick_04的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。