用pytorch搭建AlexNet(微調預訓練模型及手動搭建)
阿新 • • 發佈:2019-02-19
本文介紹瞭如何在pytorch下搭建AlexNet,使用了兩種方法,一種是直接載入預訓練模型,並根據自己的需要微調(將最後一層全連線層輸出由1000改為10),另一種是手動搭建。
構建模型類的時候需要繼承自torch.nn.Module類,要自己重寫init方法和正向傳遞時的forward方法,這裡我自己的理解是,搭建網路寫在init中,每次正向傳遞需要計算的部分寫在forward中,例如把矩陣壓平之類的。
載入預訓練alexnet之後,可以print出來檢視模型的結構及資訊:
model = models.alexnet(pretrained= True)
print(model)
分為兩個部分,features及classifier,後續搭建模型時可以也寫成這兩部分,並且從打印出來的模型資訊中也可以看出每一層的引用方式,便於修改,例如model.classifier[1]指的就是Linear(in_features=9216, out_features=4096, bias=True)這層。
下面放出完整的搭建程式碼:
import torch.nn as nn
from torchvision import models
class BuildAlexNet(nn.Module):
def __init__(self, model_type, n_output):
super(BuildAlexNet, self).__init__()
self.model_type = model_type
if model_type == 'pre':
model = models.alexnet(pretrained=True)
self.features = model.features
fc1 = nn.Linear(9216, 4096)
fc1.bias = model. classifier[1].bias
fc1.weight = model.classifier[1].weight
fc2 = nn.Linear(4096, 4096)
fc2.bias = model.classifier[4].bias
fc2.weight = model.classifier[4].weight
self.classifier = nn.Sequential(
nn.Dropout(),
fc1,
nn.ReLU(inplace=True),
nn.Dropout(),
fc2,
nn.ReLU(inplace=True),
nn.Linear(4096, n_output))
#或者直接修改為
# model.classifier[6]==nn.Linear(4096,n_output)
# self.classifier = model.classifier
if model_type == 'new':
self.features = nn.Sequential(
nn.Conv2d(3, 64, 11, 4, 2),
nn.ReLU(inplace = True),
nn.MaxPool2d(3, 2, 0),
nn.Conv2d(64, 192, 5, 1, 2),
nn.ReLU(inplace=True),
nn.MaxPool2d(3, 2, 0),
nn.Conv2d(192, 384, 3, 1, 1),
nn.ReLU(inplace = True),
nn.Conv2d(384, 256, 3, 1, 1),
nn.ReLU(inplace=True),
nn.MaxPool2d(3, 2, 0))
self.classifier = nn.Sequential(
nn.Dropout(),
nn.Linear(9216, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Linear(4096, n_output))
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
out = self.classifier(x)
return out
微調預訓練模型的思路為:直接保留原模型的features部分,重寫classifier部分。在classifier部分中,我們實際需要修改的只有最後一層全連線層,之前的兩個全連線層不需要修改,所以重寫的時候需要把這兩層的預訓練權重和偏移保留下來,也可以像註釋掉的兩行程式碼裡那樣直接引用最後一層全連線層進行修改。
網路搭好之後可以小小的測試一下以檢驗維度是否正確。
import numpy as np
from torch.autograd import Variable
import torch
if __name__ == '__main__':
model_type = 'pre'
n_output = 10
alexnet = BuildAlexNet(model_type, n_output)
print(alexnet)
x = np.random.rand(1,3,224,224)
x = x.astype(np.float32)
x_ts = torch.from_numpy(x)
x_in = Variable(x_ts)
y = alexnet(x_in)
這裡如果不加“x = x.astype(np.float32)”的話會報一個型別錯誤,感覺有點奇怪。
輸出y.data.numpy()可得10維輸出,表明網路搭建正確。