天天看點

PyTorch之六— 模型的儲存與加載、批訓練

文章目錄

  • ​​一、模型儲存與加載​​
  • ​​1.1 隻儲存和載入模型參數​​
  • ​​1.2 儲存和載入整個模型​​
  • ​​1.3 關于多GPU的模型儲存,加載等問題。​​
  • ​​二、批訓練​​
  • ​​報錯與問題解決​​

運作環境

​​

​win10 | Anaconda | PyTorch ==1.0 | python==3.6.8​

​​

這一章,主要将模型的儲存與加載,下面我們通過一個示例來示範,模型的成功儲存與加載。

一、模型儲存與加載

模型儲存和加載的兩種方式。​​關于詳情請點選​​

1.1 隻儲存和載入模型參數

這種方式需要自己定義網絡,并且其中的參數名稱與結構要與儲存的模型中的一緻(可以是部分網絡,比如隻使用VGG的前幾層),相對靈活,便于對網絡進行修改。)

PATH = "./params.pkl"
torch.save(the_model.state_dict(), PATH)     # 模型儲存

the_model = TheModelClass(*args, **kwargs)   # 初始化并載入模型
the_model.load_state_dict(torch.load(PATH))      
def save_model(model, filename):  # 若是使用GPU訓練,先将資料轉換到cpu中,再儲存
    state = model.state_dict()
    for key in state: state[key] = state[key].clone().cpu()
    torch.save(state, filename)      
1.2 儲存和載入整個模型

無需自定義網絡,儲存時已把網絡結構儲存,不能調整網絡結構。

torch.save(the_model, './model.pkl')  # 模型儲存。參數:(模型,路徑/檔案名)
the_model = torch.load('./model.pkl') # 模型加載      
1.3 關于多GPU的模型儲存,加載等問題。
  1. 若使用nn.DataParallel在一台電腦上使用了多個GPU,load模型的時候也必須先DataParallel,這和keras類似。
  2. load提供了很多重載的功能,其可以把在GPU上訓練的權重加載到CPU上跑。
model = VggNet()
model = nn.DataParallel(model).cuda()

===========================================
torch.load('tensors.pt')
# 把所有的張量加載到CPU中
torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# 把所有的張量加載到GPU 1中
torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# 把張量從GPU 1 移動到 GPU 0
torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})

# 在cpu上加載預先訓練好的GPU模型
torch.load('my_file.pt', map_location=lambda storage, loc: storage)      

上述代碼隻有在模型在一個GPU上訓練時才起作用。如果我在多個GPU上訓練我的模型,儲存它,然後嘗試在CPU上加載,我得到這個錯誤:KeyError: ‘unexpected key “module.conv1.weight” in state_dict’ 如何解決?

您可能已經使用模型儲存了模型nn.DataParallel,該模型将模型存儲在該模型中module,而現在您正試圖加載模型DataParallel。您可以nn.DataParallel在網絡中暫時添加一個加載目的,也可以加載權重檔案,建立一個沒有module字首的新的有序字典,然後加載它。

# original saved file with DataParallel
state_dict = torch.load('myfile.pth.tar') 

from collections import OrderedDict 
new_state_dict = OrderedDict()        #建立一個沒有module字首新有序字典,然後加載它。
for k, v in state_dict.items():
    name = k[7:]                      # remove `module.`
    new_state_dict[name] = v
model.load_state_dict(new_state_dict) # load params      

示例:

import torch
from  torch.autograd import Variable
import matplotlib.pyplot as plt


torch.manual_seed(1) # 随機數種子
x = torch.unsqueeze(torch.linspace(-1,1,100),dim=1)
y = x*x +0.2*torch.rand(x.size())
x,y = Variable(x,requires_grad=False),Variable(y,requires_grad=False)


def train_save():
    net1 = torch.nn.Sequential(
        torch.nn.Linear(1,10),
        torch.nn.ReLU(),
        torch.nn.Linear(10,1)
    )

    optimizer = torch.optim.SGD(net1.parameters(),lr=0.01)
    loss_func = torch.nn.MSELoss()
    
    plt.ion()
    for t in range(10000):  # 設定不同的步數100,1000
        y_pred = net1(x)
        loss = loss_func(y_pred,y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    #==================== 此段代碼用來檢視訓練過程 ==================
    #     if (t + 1) % 50 == 0:
    #         plt.cla()
    #         pred_y = y_pred.data.numpy().squeeze()
    #         plt.scatter(x.data.numpy(), y.data.numpy())
    #         plt.plot(x.data.numpy(), y_pred.data.numpy(), "r-", lw=3)
    #         print('epoch {} \t loss is {:.4f} \tacc is {:.4f}'.format(t+1,print_loss,acc))  # 訓練輪數
    #         plt.text(1.5, -4, 'acc=%.2f' % acc, fontdict={'size': 16, 'color': 'red'})
    #         plt.pause(0.05)
    # plt.ioff()
    # plt.show()
    #=============================================================
    # 繪圖
    plt.figure(1,figsize=(12,3))
    plt.subplot(131)
    plt.title("net1")
    plt.scatter(x.data.numpy(),y.data.numpy())
    plt.plot(x.data.numpy(),y_pred.data.numpy(),"r-",lw=5)

    torch.save(net1,"net1.pkl")
    torch.save(net1.state_dict(),"net_params.pkl")

def restore_net():
    net2 = torch.load("net1.pkl")
    prediction = net2(x)
    plt.subplot(132)
    plt.title("net2")
    plt.scatter(x.data.numpy(),y.data.numpy())
    plt.plot(x.data.numpy(),prediction.data.numpy(),"r-",lw=5)

def restore_params():
    net3 = torch.nn.Sequential(
        torch.nn.Linear(1,10),
        torch.nn.ReLU(),
        torch.nn.Linear(10,1)
    )
    net3.load_state_dict(torch.load("net_params.pkl"))
    prediction = net3(x)

    plt.subplot(133)
    plt.title("net3")
    plt.scatter(x.data.numpy(), y.data.numpy())
    plt.plot(x.data.numpy(), prediction.data.numpy(), "r-", lw=5)
    plt.show()

if __name__ == '__main__':
    train_save()
    restore_net()
    restore_params()      

我設定了不同的步數,展現一樣的結果。

PyTorch之六— 模型的儲存與加載、批訓練

二、批訓練

Torch 中提供了一種幫你整理你的資料結構的好東西, 叫做 DataLoader, 我們能用它來包裝自己的資料, 進行批訓練。

import torch
import torch.utils.data as Data


torch.manual_seed(1)
Batch_size = 5

x = torch.linspace(1,2,10)
y = torch.linspace(10,1,10)

# 先轉換成torch能識别的Dataset  
torch_dataset = Data.TensorDataset(x,y)

# 把dataset 放入 DataLoader
loader = Data.DataLoader(
    dataset=torch_dataset,       # torch TensorDataset資料格式
    batch_size=Batch_size,       # mini batch size
    shuffle= True,               # 是否打亂資料
    num_workers = 2              # 多線程讀取
)

if __name__ == '__main__':
    for epoch in range(3):
        for step,(batch_x,batch_y) in enumerate(loader): # 一批資料
            print('epoch:',epoch,"| step:",step,
                  "| batch_x:",batch_x.numpy(),
                  "| batch_y:", batch_y.numpy())

輸出:
epoch: 0 | step: 0 | batch_x: [1.4444444 1.6666667 2.        1.2222222 1.3333334] | batch_y: [6. 4. 1. 8. 7.]
epoch: 0 | step: 1 | batch_x: [1.1111112 1.        1.7777778 1.8888888 1.5555556] | batch_y: [ 9. 10.  3.  2.  5.]
epoch: 1 | step: 0 | batch_x: [1.3333334 1.5555556 1.6666667 2.        1.7777778] | batch_y: [7. 5. 4. 1. 3.]
epoch: 1 | step: 1 | batch_x: [1.4444444 1.2222222 1.1111112 1.        1.8888888] | batch_y: [ 6.  8.  9. 10.  2.]
epoch: 2 | step: 0 | batch_x: [1.3333334 1.1111112 1.4444444 1.5555556 2.       ] | batch_y: [7. 9. 6. 5. 1.]
epoch: 2 | step: 1 | batch_x: [1.2222222 1.8888888 1.        1.7777778 1.6666667] | batch_y: [ 8.  2. 10.  3.  4.]      
報錯與問題解決
  • 在 pytorch=0.4版本 Data.TensorDataset(data_tensor=x,target_tensor=y),現在修改為 Data.TensorDataset(x,y)
  • ​if name == ‘main‘: freeze_support()​

    ​​ 問題:在多線程環境中,mac不需要這行代碼能夠正常運作,在windows10中需要添加​

    ​if __name__ == '__main__':​

    ​才能正常運作。

第二個報錯錯誤詳情:

RuntimeError: 
        An attempt has been made to start a new process before the
        current process has finished its bootstrapping phase.

        This probably means that you are not using fork to start your
        child processes and you have forgotten to use the proper idiom
        in the main module:

            if __name__ == '__main__':
                freeze_support()
                ...

        The "freeze_support()" line can be omitted if the program
        is not going to be frozen to produce an executable.      

繼續閱讀