Pytorch預訓練模型以及修改
pytorch中自帶幾種常用的深度學習網絡預訓練模型,torchvision.models包中包含alexnet、densenet、inception、resnet、squeezenet、vgg等常用網絡結構,並且提供了預訓練模型,可通過調用來讀取網絡結構和預訓練模型(模型參數)。往往為了加快學習進度,訓練的初期直接加載pretrain模型中預先訓練好的參數。加載model如下所示:
import torchvision.models as models
1.加載網絡結構和預訓練參數:resnet34 = models.resnet34(pretrained=True)
2.#只加載網絡結構,不加載預訓練參數,即不需要用預訓練模型的參數來初始化:
resnet18 = models.resnet18(pretrained=False) #pretrained參數默認是False,為了代碼清晰,最好還是加上參數賦值.
print resnet18 #打印網絡結構
resnet18.load_state_dict(torch.load(path_params.pkl))#其中,path_params.pkl為預訓練模型參數的保存路徑。加載預先下載好的預訓練參數到resnet18,用預訓練模型的參數初始化resnet18的層,此時resnet18發生了改變。調用model的load_state_dict方法用預訓練的模型參數來初始化自己定義的新網絡結構,這個方法就是PyTorch中通用的用一個模型的參數初始化另一個模型的層的操作。load_state_dict方法還有一個重要的參數是strict,該參數默認是True,表示預訓練模型的層和自己定義的網絡結構層嚴格對應相等(比如層名和維度)。故,當新定義的網絡(model_dict)和預訓練網絡(pretrained_dict)的層名不嚴格相等時,需要先將pretrained_dict里不屬於model_dict的鍵剔除掉 :
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} ,再用預訓練模型參數更新model_dict,最后用load_state_dict方法初始化自己定義的新網絡結構。
print resnet18 #打印的還是網絡結構
注意: cnn = resnet18.load_state_dict(torch.load( path_params.pkl )) #是錯誤的,這樣cnn將是nonetype
pre_dict = resnet18.state_dict() #按鍵值對將模型參數加載到pre_dict
print for k, v in pre_dict.items(): # 打印模型參數
for k, v in pre_dict.items():
print k #打印模型每層命名
#model是自己定義好的新網絡模型,將pretrained_dict和model_dict中命名一致的層加入pretrained_dict(包括參數)。
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
預訓練模型的修改(具體要求不同,則用到的修改方式不同。)
1. 參數修改
對於簡單的參數修改,這里以resnet預訓練模型舉例,resnet源代碼在Github。 resnet網絡最后一層分類層fc是對1000種類型進行划分,對於自己的數據集,如果只有9類,修改的代碼如下:
# coding=UTF-8
import torchvision.models as models
#調用模型
model = models.resnet50(pretrained=True)
#提取fc層中固定的參數
fc_features = model.fc.in_features
#修改類別為9
model.fc = nn.Linear(fc_features, 9)
2. 增減卷積層 (待補充)
前一種方法只適用於簡單的參數修改,有時候往往要修改網絡中的層次結構,這時只能用參數覆蓋的方法,即自己先定義一個類似的網絡,再將預訓練中的參數提取到自己的網絡中來。這里以resnet預訓練模型舉例。
3. 訓練特定層,凍結其它層
另一種使用預訓練模型的方法是對它進行部分訓練。具體做法是,將模型起始的一些層的權重保持不變,重新訓練后面的層,得到新的權重。在這個過程中,可多次進行嘗試,從而能夠依據結果找到frozen layers和retrain layers之間的最佳搭配。
如何使用預訓練模型,是由數據集大小和新舊數據集(預訓練的數據集和自己要解決的數據集)之間數據的相似度來決定的。
下圖表展示了在各種情況下應該如何使用預訓練模型: 
Pytorch保存與加載網絡模型,有以下兩種方式: (推薦第二種)
一、是保存整個神經網絡的的結構信息和模型參數信息,save的對象是網絡net;
torch.save(model_object, 'model.pkl') # 保存整個神經網絡的結構和模型參數
重載:model = torch.load('model.pkl') #重載並初始化新的神經網絡對象。
二、是只保存神經網絡的訓練模型參數,save的對象是net.state_dict()。
torch.save(model_object.state_dict(), 'params.pkl') # 只保存神經網絡的模型參數
需要首先導入對應的網絡,通過model_object.load_state_dict(torch.load('params.pkl'))完成模型參數的重載和初始化新定義的網絡。
PyTorch中使用預訓練的模型初始化網絡的一部分參數
#首先自己新定義一個網絡
class CNN(nn.Module):
def __init__(self, block, layers, num_classes=9): #自己新定義的CNN與繼承的ResNet網絡結構大體相同,即除了新增層,其他層的層名與ResNet的相同。
self.inplanes = 64
super(ResNet, self).__init__() #繼承ResNet網絡結構
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
self.avgpool = nn.AvgPool2d(7, stride=1)
#新增一個反卷積層
self.convtranspose1 = nn.ConvTranspose2d(2048, 2048, kernel_size=3, stride=1, padding=1, output_padding=0, groups=1, bias=False, dilation=1)
#新增一個最大池化層
self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
#將原來的fc層改成fclass層
self.fclass = nn.Linear(2048, num_classes) #原來的fc層:self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules(): #
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels #
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(self.inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes * block.expansion),
)
layers = [ ]
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
#3個新加層的forward
x = x.view(x.size(0), -1) #因為接下來的self.convtranspose1層的輸入通道是2048
x = self.convtranspose1(x)
x = self.maxpool2(x)
x = x.view(x.size(0), -1) #因為接下來的self.fclass層的輸入通道是2048
x = self.fclass(x)
return x
#加載model
resnet50 = models.resnet50(pretrained=True)
cnn = CNN(Bottleneck, [3, 4, 6, 3]) #創建一個自己新定義的網絡對象cnn。
pretrained_dict = resnet50.state_dict() #記錄預訓練模型的參數:resnet50.state_dict()。若已存在 resnet50.state_dict()對應的模型參數文件 'params.pkl',則此句代碼等價於:pretrained_dict =torch.load(path_params.pkl) ?其中,path_params.pkl為' params.pkl '的保存路徑。
model_dict = cnn.state_dict() #自己新定義網絡的參數
# 將pretrained_dict里不屬於model_dict的鍵剔除掉 ,因為后面的cnn.load_state_dict()方法有個重要參數是strict,默認是True,表示預訓練模型的層和自己定義的網絡結構層嚴格對應相等(比如層名和維度)。
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} #只能對層名一致的層進行“層名:參數”鍵值對賦值。
# 更新現有的model_dict
model_dict.update(pretrained_dict)
# 加載我們真正需要的模型參數state_dict
cnn.load_state_dict(model_dict) #cnn.load_state_dict()方法對cnn初始化,其一個重要參數strict,默認為True,表示預訓練模型(model_dict)的層和自己定義的網絡結構(cnn)的層嚴格對應相等(比如層名和維度)。
print(cnn)
后續在此基礎上繼續重新進行訓練,如下面即將介紹的: 選擇特定的層進行finetune 。
選擇特定的層進行finetune
先使用Module.children()方法查看網絡的直接子模塊,將不需要調整的模塊中的參數設置為param.requires_grad = False,同時用一個list收集需要調整的模塊中的參數。具體代碼為:
count = 0
para_optim = []
for k in model.children():
count += 1
if count > 6: # 6 should be changed properly
for param in k.parameters():
para_optim.append(param)
else:
for param in k.parameters():
param.requires_grad = False
optimizer = optim.RMSprop(para_optim, lr)#只對特定的層的參數進行優化更新,即選擇特定的層進行finetune。
到此我們實現了PyTorch中使用預訓練的模型初始化網絡的一部分參數。
此部分主要參考PyTorch教程的Autograd machnics部分
1.在PyTorch中,每個Variable數據含有兩個flag(requires_grad和volatile)用於指示是否計算此Variable的梯度。設置requires_grad = False,或者設置volatile=True,即可指示不計算此Variable的梯度:
for param in model.parameters():
param.requires_grad = False
注意,在模型測試時,對input_data設置volatile=True,可以節省測試時的顯存 。
2.PyTorch的Module.modules()和Module.children()
參考PyTorch document和discuss
在PyTorch中,所有的neural network module都是class torch.nn.Module的子類,在Modules中可以包含其它的Modules,以一種樹狀結構進行嵌套。當需要返回神經網絡中的各個模塊時,Module.modules()方法返回網絡中所有模塊的一個iterator,而Module.children()方法返回所有直接子模塊的一個iterator。具體而言:
list ( nn.Sequential(nn.Linear(10, 20), nn.ReLU()).modules() )
Out[9]:
[Sequential (
(0): Linear (10 -> 20)
(1): ReLU ()
), Linear (10 -> 20), ReLU ()]
In [10]: list( nn.Sequential(nn.Linear(10, 20), nn.ReLU()) .children() )
Out[10]: [Linear (10 -> 20), ReLU ()]
舉例:Faster-RCNN基於vgg19提取features,但是只使用了vgg19一部分模型提取features。
步驟:
下載vgg19的pth文件,在anaconda中直接設置pretrained=True下載一般都比較慢,在model_zoo里面有各種預訓練模型的下載鏈接:
model_urls = {
'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
'vgg11_bn': 'https://download.pytorch.org/models/vgg11_bn-6002323d.pth',
'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth',
'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth' }
下載好的模型,可以用下面這段代碼看一下模型參數,並且改一下模型。在vgg19.pth同級目錄建立一個test.py。
import torch
import torch.nn as nn
import torchvision.models as models
vgg16 = models.vgg16(pretrained=False)
#打印出預訓練模型的參數
vgg16.load_state_dict(torch.load('vgg16-397923af.pth'))
print('vgg16:\n', vgg16)
modified_features = nn.Sequential(*list(vgg16.features.children())[:-1])
# to relu5_3
print('modified_features:\n', modified_features )#打印修改后的模型參數
修改好之后features就可以拿去做Faster-RCNN提取特征用了。
