batchnorm2d參數 torch_Pytorch自由載入部分模型參數并凍結的示例分析

這篇文章給大家介紹batchnorm2d參數 torch_Pytorch自由載入部分模型參數并凍結的示例分析，內容非常詳細，感興趣的小伙伴們可以參考借鑒，希望對大家能有所幫助。

Pytorch的load方法和load_state_dict方法只能較為固定的讀入參數文件，他們要求讀入的state_dict的key和Model.state_dict()的key對應相等。

而我們在進行遷移學習的過程中也許只需要使用某個預訓練網絡的一部分，把多個網絡拼和成一個網絡，或者為了得到中間層的輸出而分離預訓練模型中的Sequential 等等，這些情況下。傳統的load方法就不是很有效了。

例如，我們想利用Mobilenet的前7個卷積并把這幾層凍結，后面的部分接別的結構，或者改寫成FCN結構，傳統的方法就不奏效了。

最普適的方法是：構建一個字典，使得字典的keys和我們自己創建的網絡相同，我們再從各種預訓練網絡把想要的參數對著新的keys填進去就可以有一個新的state_dict了，這樣我們就可以load這個新的state_dict，目前只能想到這個方法應對較為復雜的網絡變換。

網上查“載入部分模型”，“凍結部分模型”一般都是只改個FC,根本沒有用，初學的時候自己寫state_dict也踩了一些坑，發出來記錄一下。

一.載入部分預訓練參數

我們先看看Mobilenet的結構

( 來源github,附帶預訓練模型mobilenet_sgd_rmsprop_69.526.tar）

class Net(nn.Module):

def __init__(self):

super(Net, self).__init__()


def conv_bn(inp, oup, stride):

return nn.Sequential(

nn.Conv2d(inp, oup, 3, stride, 1, bias=False),

nn.BatchNorm2d(oup),

nn.ReLU(inplace=True)

)


def conv_dw(inp, oup, stride):

return nn.Sequential(

nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),

nn.BatchNorm2d(inp),

nn.ReLU(inplace=True),


nn.Conv2d(inp, oup, 1, 1, 0, bias=False),

nn.BatchNorm2d(oup),

nn.ReLU(inplace=True),

)


self.model = nn.Sequential(

conv_bn( 3, 32, 2),

conv_dw( 32, 64, 1),

conv_dw( 64, 128, 2),

conv_dw(128, 128, 1),

conv_dw(128, 256, 2),

conv_dw(256, 256, 1),

conv_dw(256, 512, 2),

conv_dw(512, 512, 1),

conv_dw(512, 512, 1),

conv_dw(512, 512, 1),

conv_dw(512, 512, 1),

conv_dw(512, 512, 1),

conv_dw(512, 1024, 2),

conv_dw(1024, 1024, 1),

nn.AvgPool2d(7),

)

self.fc = nn.Linear(1024, 1000)


def forward(self, x):

x = self.model(x)

x = x.view(-1, 1024)

x = self.fc(x)

return x

我們只需要前7層卷積，并且為了方便日后concate操作，我們把Sequential拆開，成為下面的樣子

class Net(nn.Module):

def __init__(self):

super(Net, self).__init__()


def conv_bn(inp, oup, stride):

return nn.Sequential(

nn.Conv2d(inp, oup, 3, stride, 1, bias=False),

nn.BatchNorm2d(oup),

nn.ReLU(inplace=True)

)


def conv_dw(inp, oup, stride):

return nn.Sequential(

nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),

nn.BatchNorm2d(inp),

nn.ReLU(inplace=True),


nn.Conv2d(inp, oup, 1, 1, 0, bias=False),

nn.BatchNorm2d(oup),

nn.ReLU(inplace=True),

)


self.conv1 = conv_bn( 3, 32, 2)

self.conv2 = conv_dw( 32, 64, 1)

self.conv3 = conv_dw( 64, 128, 2)

self.conv4 = conv_dw(128, 128, 1)

self.conv5 = conv_dw(128, 256, 2)

self.conv6 = conv_dw(256, 256, 1)

self.conv7 = conv_dw(256, 512, 2)


# 原來這些不要了

# 可以自己接后面的結構

'''

self.features = nn.Sequential(

conv_dw(512, 512, 1),

conv_dw(512, 512, 1),

conv_dw(512, 512, 1),

conv_dw(512, 512, 1),

conv_dw(512, 512, 1),

conv_dw(512, 1024, 2),

conv_dw(1024, 1024, 1),

nn.AvgPool2d(7),)


self.fc = nn.Linear(1024, 1000)

'''


def forward(self, x):

x1 = self.conv1(x)

x2 = self.conv2(x1)

x3 = self.conv3(x2)

x4 = self.conv4(x3)

x5 = self.conv5(x4)

x6 = self.conv6(x5)

x7 = self.conv7(x6)

#x8 = self.features(x7)

#out = self.fc

return (x1,x2,x3,x4,x4,x6,x7)

我們更具改過的結構創建一個net,看看他的state_dict和我們預訓練文件的state_dict有啥區別

net = Net()

#我的電腦沒有GPU,他的參數是GPU訓練的cudatensor,于是要下面這樣轉換一下

dict_trained = torch.load("mobilenet_sgd_rmsprop_69.526.tar",map_location=lambda storage, loc: storage)["state_dict"]

dict_new = net.state_dict().copy()


new_list = list (net.state_dict().keys() )

trained_list = list (dict_trained.keys() )

print("new_state_dict size: {} trained state_dict size: {}".format(len(new_list),len(trained_list)) )

print("New state_dict first 10th parameters names")

print(new_list[:10])

print("trained state_dict first 10th parameters names")

print(trained_list[:10])


print(type(dict_new))

print(type(dict_trained))

得到輸出如下：

我們截斷一半之后，參數由137變成65了，前十個參數看出，名字變了但是順序其實沒變。state_dict的數據類型是Odict,可以按照dict的操作方法操作。

new_state_dict size: 65 trained state_dict size: 137
New state_dict first 10th parameters names
['conv1.0.weight', 'conv1.1.weight', 'conv1.1.bias', 'conv1.1.running_mean', 'conv1.1.running_var', 'conv2.0.weight', 'conv2.1.weight', 'conv2.1.bias', 'conv2.1.running_mean', 'conv2.1.running_var']
trained state_dict first 10th parameters names
['module.model.0.0.weight', 'module.model.0.1.weight', 'module.model.0.1.bias', 'module.model.0.1.running_mean', 'module.model.0.1.running_var', 'module.model.1.0.weight', 'module.model.1.1.weight', 'module.model.1.1.bias', 'module.model.1.1.running_mean', 'module.model.1.1.running_var']
<class 'collections.OrderedDict'>
<class 'collections.OrderedDict'>

我們看出只要構建一個字典，使得字典的keys和我們自己創建的網絡相同，我們在從各種預訓練網絡把想要的參數對著新的keys填進去就可以有一個新的state_dict了，這樣我們就可以load這個新的state_dict，這是最普適的方法適用于所有的網絡變化。

for i in range(65):

dict_new[ new_list[i] ] = dict_trained[ trained_list[i] ]


net.load_state_dict(dict_new)

還有別的情況，比如我們只是在后面加了一些層，沒有改變原來網絡層的名字和結構，可以用下面的簡便方法：

loaded_dict = {k: loaded_dict[k] for k, _ in model.state_dict()}

二.凍結這幾層參數

方法很多，這里用和上面方法對應的凍結方法

發現之前的凍結有問題，還是建議看一下

https://discuss.pytorch.org/t/how-the-pytorch-freeze-network-in-some-layers-only-the-rest-of-the-training/7088

或者

https://discuss.pytorch.org/t/correct-way-to-freeze-layers/26714

或者

對應的，在訓練時候，optimizer里面只能更新requires_grad = True的參數，于是

optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, net.parameters(),lr) )

關于batchnorm2d參數 torch_Pytorch自由載入部分模型參數并凍結的示例分析就分享到這里了，希望以上內容可以對大家有一定的幫助，可以學到更多知識。如果覺得文章不錯，可以把它分享出去讓更多的人看到。