pytorch在fintune時將sequential中的層輸出,以vgg為例
阿新 • • 發佈:2019-01-07
pytorch將sequential中的層輸出,以vgg為例
有時候我們在fintune時發現pytorch把許多層都集合在一個sequential裡,但是我們希望能把中間層的結果引出來做下一步操作,於是我自己琢磨了一個方法,以vgg為例,有點僵硬哈!
首先pytorch自帶的vgg16模型的網路結構如下:
VGG(
(features): Sequential(
(0): Conv2d (3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace)
(2): Conv2d (64 , 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): ReLU(inplace)
(4): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1))
(5): Conv2d (64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(6): ReLU(inplace)
(7): Conv2d (128, 128, kernel_size=(3, 3), stride=(1 , 1), padding=(1, 1))
(8): ReLU(inplace)
(9): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1))
(10): Conv2d (128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): ReLU(inplace)
(12): Conv2d (256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(13 ): ReLU(inplace)
(14): Conv2d (256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(15): ReLU(inplace)
(16): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1))
(17): Conv2d (256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(18): ReLU(inplace)
(19): Conv2d (512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(20): ReLU(inplace)
(21): Conv2d (512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(22): ReLU(inplace)
(23): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1))
(24): Conv2d (512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(25): ReLU(inplace)
(26): Conv2d (512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(27): ReLU(inplace)
(28): Conv2d (512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(29): ReLU(inplace)
(30): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1))
)
(classifier): Sequential(
(0): Linear(in_features=25088, out_features=4096)
(1): ReLU(inplace)
(2): Dropout(p=0.5)
(3): Linear(in_features=4096, out_features=4096)
(4): ReLU(inplace)
(5): Dropout(p=0.5)
(6): Linear(in_features=4096, out_features=1000)
)
)
我們需要fintune vgg16的features部分,並且我希望把3,8, 15, 22, 29這五個作為輸出進一步操作。我的想法是自己寫一個vgg網路,這個網路引數與pytorch的網路一致但是保證我們需要的層輸出在sequential外。於是我寫的網路如下:
class our_vgg(nn.Module):
def __init__(self):
super(our_vgg, self).__init__()
self.conv1 = nn.Sequential(
# conv1
nn.Conv2d(3, 64, 3, padding=35),
nn.ReLU(inplace=True),
nn.Conv2d(64, 64, 3, padding=1),
nn.ReLU(inplace=True),
)
self.conv2 = nn.Sequential(
# conv2
nn.MaxPool2d(2, stride=2, ceil_mode=True), # 1/2
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(128, 128, 3, padding=1),
nn.ReLU(inplace=True),
)
self.conv3 = nn.Sequential(
# conv3
nn.MaxPool2d(2, stride=2, ceil_mode=True), # 1/4
nn.Conv2d(128, 256, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, 3, padding=1),
nn.ReLU(inplace=True),
)
self.conv4 = nn.Sequential(
# conv4
nn.MaxPool2d(2, stride=2, ceil_mode=True), # 1/8
nn.Conv2d(256, 512, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, 3, padding=1),
nn.ReLU(inplace=True),
)
self.conv5 = nn.Sequential(
# conv5
nn.MaxPool2d(2, stride=2, ceil_mode=True), # 1/16
nn.Conv2d(512, 512, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, 3, padding=1),
nn.ReLU(inplace=True),
)
def forward(self, x):
conv1 = self.conv1(x)
conv2 = self.conv2(conv1)
conv3 = self.conv3(conv2)
conv4 = self.conv4(conv3)
conv5 = self.conv5(conv4)
return conv5
接著就是copy weights了:
def convert_vgg(vgg16):#vgg16是pytorch自帶的
net = our_vgg()# 我寫的vgg
vgg_items = net.state_dict().items()
vgg16_items = vgg16.items()
pretrain_model = {}
j = 0
for k, v in net.state_dict().iteritems():#按順序依次填入
v = vgg16_items[j][1]
k = vgg_items[j][0]
pretrain_model[k] = v
j += 1
return pretrain_model
## net是我們最後使用的網路,也是我們想要放置weights的網路
net = net()
print ('load the weight from vgg')
pretrained_dict = torch.load('vgg16.pth')
pretrained_dict = convert_vgg(pretrained_dict)
model_dict = net.state_dict()
# 1. 把不屬於我們需要的層剔除
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. 把引數存入已經存在的model_dict
model_dict.update(pretrained_dict)
# 3. 載入更新後的model_dict
net.load_state_dict(model_dict)
print ('copy the weight sucessfully')
這樣我就基本達成目標了,注意net也就是我們要使用的網路fintune部分需要和our_vgg一致。