轉載：https://zhuanlan.zhihu.com/p/53927068
https://blog.csdn.net/wangdongwei0/article/details/88956527
pytorch最后的權重文件是.pth格式的。
經常遇到的問題：
進行finutune時，改配置文件中的學習率，發現程序跑起來后竟然保持了以前的學習率， 並沒有使用新的學習率。
原因：
首先查看.pth文件中的內容，我們發現它其實是一個字典格式的文件：
其中保存了optimizer和scheduler，所以再次加載此文件時會使用之前的學習率。
我們只需要權重，也就是model部分，將其導出就可以了
import torch original = torch.load('path/to/your/checkpoint.pth') new = {"model": original["model"]} torch.save(new, 'path/to/new/checkpoint.pth')
背景
在PyTroch框架中，如果要自定義一個Net(網絡，或者model，在本文中，model和Net擁有同樣的意思)，通常需要繼承自nn.Module然后實現自己的layer。比如，在下面的示例中，gemfield（tiande亦有貢獻）使用Pytorch實現了一個Net（可以看到其父類為nn.Module)：
import torch import torch.nn as nn import torch.nn.functional as F class CivilNet(nn.Module): def __init__(self): super(CivilNet, self).__init__() self.conv1 = nn.Conv2d(3, 6, 5) self.pool = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(6, 16, 5) self.fc1 = nn.Linear(16 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) self.gemfield = "gemfield.org" self.syszux = torch.zeros([1,1]) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = x.view(-1, 16 * 5 * 5) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x
這就帶來了一系列的問題：
1，為什么要繼承自nn.Module？
2，網絡的各個layer或者module為什么要直接定義在構造函數中，而不能（比方說）放在構造函數中的一個list里？
3，forward函數什么時候會被調用？為什么要使用net(input)而不是net.forward(input)來做前向呢？
4，保存模型時，保存的究竟是什么？
5，重新載入一個pth模型時，究竟發生了什么？
你肯定要問了，為什么沒說到反向？因為反向是optimizer和tensor的grad共同完成的，本文只討論Net部分，這一系列文章的后續部分會討論反向。
CivilNet的實例化
一個Net，也就是繼承自nn.Module的類，當實例化后，本質上就是維護了以下8個字典(OrderedDict)：
_parameters
_buffers
_backward_hooks
_forward_hooks
_forward_pre_hooks
_state_dict_hooks
_load_state_dict_pre_hooks
_modules
這8個字典用於網絡的前向、反向、序列化、反序列化中。
因此，當實例化你定義的Net(nn.Module的子類)時，要確保父類的構造函數首先被調用，這樣才能確保上述8個OrderedDict被create出來，否則，后續任何的初始化操作將拋出類似這樣的異常：cannot assign module before Module.__init__() call。
對於前述的CivilNet而言，當CivilNet被實例化后，CivilNet本身維護了這8個OrderedDict，更重要的是，CivilNet中的conv1和conv2(類型為nn.modules.conv.Conv2d）、pool（類型為nn.modules.pooling.MaxPool2d）、fc1、fc2、fc3（類型為torch.nn.modules.linear.Linear）均維護了8個OrderedDict，因為它們的父類都是nn.Module，而gemfield（類型為str）、syszux（類型為torch.Tensor)則沒有這8個OrderedDict。
也因此，在你定義的網絡投入運行前，必然要確保和上面一樣——構造出那8個OrderedDict，這個構造，就在nn.Module的構造函數中。如此以來，你定義的Net就必須繼承自nn.Module；如果你的Net定義了__init__()方法，則必須在你的__init__方法中調用nn.Module的構造函數，比如super(your_class).__init__() ，注意，如果你的子類沒有定義__init__()方法，則在實例化的時候會默認用nn.Module的，這種情況也對。
nn.Module通過使用__setattr__機制，使得定義在類中（不一定要定義在構造函數里）的成員（比如各種layer），被有序歸屬到_parameters、_modules、_buffers或者普通的attribute里；那具體怎么歸屬呢？很簡單，當類成員的type 派生於Parameter類時（比如conv的weight，在CivilNet類中，就是self.conv1中的weight屬性），該屬性就會被划歸為_parameters；當類成員的type派生於Module時（比如CivilNet中的self.conv1，其實除了gemfield和syszux外都是），該成員就會划歸為_modules。
如果知道了這個機制，就會自然而然的知道，如果上面的CivilNet里的成員封裝到一個list里，像下面這樣：
class CivilNet(nn.Module): def __init__(self): super(CivilNet, self).__init__() conv1 = nn.Conv2d(3, 6, 5) pool = nn.MaxPool2d(2, 2) conv2 = nn.Conv2d(6, 16, 5) self.layer1 = [conv1, pool, conv2] ...
那么在運行的時候，可能optimizer就會提示parameters為empty。這就是因為成員layer1的type派生自list，而非Module；而像CivilNet這樣的Net，在取所有的parameters的時候，都是通過_modules橋梁去取得的......
1，_parameters
前述說到了parameters就是Net的權重參數（比如conv的weight、conv的bias、fc的weight、fc的bias)，類型為tensor，用於前向和反向；比如，你針對Net使用cpu()、cuda()等調用的時候，實際上調用的就是parameter這個tensor的cpu()、cuda()等方法；再比如，你保存模型或者重新加載pth文件的時候，針對的都是parameter的操作或者賦值。
如果你針對的是CivilNet直接取_parameters屬性的值的話，很遺憾是空的，因為CivilNet的成員並沒有直接派生自Parameter類；但是當針對CivilNet取parameters()函數的返回值（是個iter）時，則會遞歸拿到所有的，比如conv的weight、bias等；
2，_buffers
該成員值的填充是通過register_buffer API來完成的，通常用來將一些需要持久化的狀態（但又不是網絡的參數）放到_buffer里；一些極其個別的操作，比如BN，會將running_mean的值放入進來；
3，_modules
_modules成員起很重要的橋梁作用，在獲取一個net的所有的parameters的時候，是通過遞歸遍歷該net的所有_modules來實現的。
像前述提到的那個問題，如果將這些成員都放倒一個python list里：self.layer1 = [conv1, pool, conv2] ——會導致CivilNet不能將conv1, pool, conv2等划歸到_modules里，從而通過CivilNet的parameters()獲取所有權重參數時，拿到的東西為空，就會報optimizer got an empty parameter list這樣的錯誤。針對這種情況，那怎么辦呢？
ModuleList就是為了解決這個問題的，首先，ModuleList類的基類正是Module：
class ModuleList(Module)
其次，ModuleList實現了python的list的功能；
最后，在使用ModuleList的時候，該類會使用基類（也就是Module）的add_module()方法，或者直接操作_modules成員來將list中的module成功注冊。
Sequential模塊也具備ModuleList這樣的注冊功能，另外其還實現了forward，這是和ModuleList不同的地方：
def forward(self, input): for module in self._modules.values(): input = module(input) return input
CivilNet的前向
網絡的前向需要通過諸如CivilNet(input)這樣的形式來調用，而非CivilNet.forward(input)，是因為前者實現了額外的功能：
1，先執行完所有的_forward_pre_hooks里的hooks 2, 再調用CivilNet的forward函數 3, 再執行完所有的_forward_hooks中的hooks 4, 再執行完所有的_backward_hooks中的hooks
可以看到:
1，_forward_pre_hooks是在網絡的forward之前執行的。這些hooks通過網絡的register_forward_pre_hook() API來完成注冊，通常只有一些Norm操作會定義_forward_pre_hooks。這種hook不能改變input的內容。
2，_forward_hooks是通過register_forward_hook來完成注冊的。這些hooks是在forward完之后被調用的，並且不應該改變input和output。目前就是方便自己測試的時候可以用下。
3，_backward_hooks和_forward_hooks類似。
所以總結起來就是，如果你的網絡中沒有Norm操作，那么使用CivilNet(input)和CivilNet.forward(input)是等價的。
另外，你必須使用CivilNet.eval()操作來將dropout和BN這些op設置為eval模式，否則你將得到不一致的前向返回值。eval()調用會將Net的實例中的training成員設置為False。
CivilNet模型的保存和重新加載
如果我們要保存一個訓練好哦PyTorch模型的話，會使用下面的API：
cn = CivilNet() ...... torch.save(cn.state_dict(), "your_model_path.pth")
可以看到使用了網絡的state_dict() API調用以及torch模塊的save調用。一言以蔽之，模型的保存就是先通過state_dict() API的調用獲得一個關於網絡參數的字典，再通過pickle模塊序列化成文件的形式。
而如果我們要load一個pth模型來進行前向的時候，會使用下面的API：
cn = CivilNet() #參數反序列化為python dict state_dict = torch.load("your_model_path.pth") #加載訓練好的參數 cn.load_state_dict(state_dict) #變成測試模式，dropout和BN在訓練和測試時不一樣 #eval()會把模型中的每個module的self.training設置為False cn = cn.cuda().eval()
可以看到使用了torch模塊的load調用和網絡的load_state_dict() API調用。一言以蔽之，模型的重新加載就是先通過torch.load反序列化pickle文件得到一個Dict，然后再使用該Dict去初始化當前網絡的state_dict。torch的save和load API在python2中使用的是cPickle，在python3中使用的是pickle。另外需要注意的是，序列化的pth文件會被寫入header信息，包括magic number、version信息等。
關於模型的保存，我們需要弄清楚以下概念：1, state_dict；2, 序列化一個pth模型用於以后的前向；3, 為之后的再訓練保存一個中間的checkpoint；4,將多個模型保存為一個文件；5,用其它模型的參數來初始化當前的網絡；6,跨設備的模型的保存和加載。
1, state_dict
在Pytorch中，可學習的參數(如Module中的weights和biases)是包含在網絡的parameters()調用返回的字典中的，這就是一個普通的OrderedDict，這里面的key-value是通過網絡及遞歸網絡里的Module成員獲取到的：它的key是每一個layer的成員的名字(加上prefix），而對應的value是一個tensor。比如本文前述的CivilNet類，它的state_dict中的key如下所示：
conv1.weight
conv1.bias
conv2.weight
conv2.bias
fc1.weight
fc1.bias
fc2.weight
fc2.bias
fc3.weight
fc3.bias
那如果你使用了DataParallel來訓練的話：
cn = nn.DataParallel(cn)
那么state_dict中的key將如下所示：
module.conv1.weight
module.conv1.bias
module.conv2.weight
module.conv2.bias
module.fc1.weight
module.fc1.bias
module.fc2.weight
module.fc2.bias
module.fc3.weight
module.fc3.bias
如果你使用了ModuleList的話，比如前述CivilNet的定義你寫作了：
class CivilNet(nn.Module): def __init__(self): super(CivilNet, self).__init__() conv1 = nn.Conv2d(3, 6, 5) pool = nn.MaxPool2d(2, 2) conv2 = nn.Conv2d(6, 16, 5) fc1 = nn.Linear(16 * 5 * 5, 120) fc2 = nn.Linear(120, 84) fc3 = nn.Linear(84, 10) self.gemfield = nn.ModuleList([conv1, pool, conv2, fc1, fc2, fc3])
那state_dict中的key將如下所示：
gemfield.1.weight gemfield.1.bias gemfield.2.weight gemfield.2.bias gemfield.3.weight gemfield.3.bias gemfield.4.weight gemfield.4.bias gemfield.5.weight gemfield.5.bias
還有很多的變種，不過大抵上你也知道規律了。
2，load_state_dict
load_state_dict()調用是nn.Module的一個API，用模型文件反序列化后得到的Dict來初始化當前的模型。需要提及的是這個函數上的 strict參數，默認值是True。因此在初始化時候，該函數會嚴格比較源Dict和目標Dict的key是否一樣，不能多也不能少，必須嚴格一樣。
如果將strict參數設置為False，則將不會進行這樣嚴格的check。只有key一樣的才會進行賦值。
3,序列化模型以保存state_dict
這種情況是PyTorch中最常用的保存模型的方法。
#save torch.save(model.state_dict(), PATH) #load model = CivilNet(*args, **kwargs) model.load_state_dict(torch.load(PATH)) model.eval()
不再贅述。
4，序列化整個模型
#save torch.save(model, PATH) #load model = torch.load(PATH) model.eval()
這種方式不推薦，其是通過Pickle模塊將整個class序列化了，序列化過程中依賴很多具體的東西，比如定義model class的路徑。這樣反序列化的時候就喪失了靈活性。
5，序列化中間過程中的checkpoint
這種序列化的目的是為了之后以這個狀態為基點重新開始訓練。和前述序列化模型的本質不同就在於還需要序列化optimizer的Dict（比如學習率等參數）。傳統上，checkpoint文件用.tar作為后綴：
#save torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': loss, ... }, PATH) #load model = CivilNet(*args, **kwargs) optimizer = TheOptimizerClass(*args, **kwargs) checkpoint = torch.load(PATH) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) epoch = checkpoint['epoch'] loss = checkpoint['loss'] model.train() #model.eval()
6，將多個模型序列化到一個文件里
比如，decoder-encoder這種結構會有多個Net。傳統上，checkpoint文件用.tar作為后綴。
#save torch.save({ 'modelA_state_dict': modelA.state_dict(), 'modelB_state_dict': modelB.state_dict(), 'optimizerA_state_dict': optimizerA.state_dict(), 'optimizerB_state_dict': optimizerB.state_dict(), ... }, PATH) #load modelA = TheModelAClass(*args, **kwargs) modelB = TheModelBClass(*args, **kwargs) optimizerA = TheOptimizerAClass(*args, **kwargs) optimizerB = TheOptimizerBClass(*args, **kwargs) checkpoint = torch.load(PATH) modelA.load_state_dict(checkpoint['modelA_state_dict']) modelB.load_state_dict(checkpoint['modelB_state_dict']) optimizerA.load_state_dict(checkpoint['optimizerA_state_dict']) optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])
7，用一個模型的部分參數初始化另一個模型（遷移學習）
這種情況的目的是為了復用一個模型的部分layer，以實現遷移學習。
#save torch.save(modelA.state_dict(), PATH) #load modelB = TheModelBClass(*args, **kwargs) modelB.load_state_dict(torch.load(PATH), strict=False)
和前述保存模型相比，序列化部分一樣，反序列化只需要將strict參數設置為False。在前述load_state_dict章節中已經解釋過，此處不再贅述。
8，跨device（cpu/gpu）來save/load模型
比如模型是在GPU上訓練的，現在要load到cpu上。或者反之，或者在CPU上訓練，在GPU上load。這三種情況下，save的方法是一樣的：
torch.save(model.state_dict(), PATH)
而load的方法就不一樣了：
###############Save on GPU, Load on CPU ######### device = torch.device('cpu') model = CivilNet(*args, **kwargs) model.load_state_dict(torch.load(PATH, map_location=device)) ###############Save on GPU, Load on GPU ######### device = torch.device("cuda") model = CivilNet(*args, **kwargs) model.load_state_dict(torch.load(PATH)) model.to(device) #確保在輸入給網絡的tensor上調用input = input.to(device) ###############Save on CPU, Load on GPU ######### device = torch.device("cuda") model = CivilNet(*args, **kwargs) model.load_state_dict(torch.load(PATH, map_location="cuda:0")) # Choose whatever GPU device number you want model.to(device) #確保在輸入給網絡的tensor上調用input = input.to(device)
9，使用torch.nn.DataParallel訓練的模型如何序列化
torch.nn.DataParallel 是一個wrapper，用來幫助在多個GPU上並行進行運算。這種情況下要保存訓練好的模型，最好使用model.module.state_dict()，請參考本章第1節：state_dict。這種情況下你在重新加載pth模型文件的時候，就會有極大的靈活性，而不是出現一大堆unexpected keys和missed keys：
torch.save(model.module.state_dict(), PATH)
打印CivilNet
這個是靠__repr__機制，不再贅述；
cn = CivilNet()
print(cn)
另外，你的類可以重寫nn.Module的extra_repr()方法來實現定制化的打印。