AlexNet-Pytorch-Kaggle貓狗大戰
前言
前一段時間基於LeNet-5實現了MNIST手寫數字識別,由於torchvision.datasets模組集成了MNIST資料集,所以在載入資料時使用的是torchvision.datasets自帶的方法,缺失瞭如何對一般資料集的處理部分,不能將其作為一個模板來適用於新的網路。通常,我們需要為待處理的資料集定義一個單獨的資料處理類,在本文中,將基於AlexNet來實現貓狗分類,並詳細總結各個部分。對於我自己來說,在後面適用新的網路時,希望能夠以此次的程式碼作為一個模組,增加效率,這也是寫這篇部落格的目的所在。
相關資料下載地址
AlexNet論文地址:https://papers.nips.cc/paper/4824-imagenet-classification-with-deep-convolutional-neural-networks.pdf)
專案地址:https://github.com/myCigar/cat_vs_dog-AlexNet
資料下載地址:https://www.kaggle.com/c/dogs-vs-cats-redux-kernels-edition/data
AlexNet網路構建
AlexNet由Hinton和他的學生Alex Krizhevsky在2012年所提出,並在當年的ImageNet競賽中獲得了冠軍,在論文中還提出了ReLu,Dropout,LRN等用於優化網路的功能,ReLu啟用函式加快了訓練的速度,Dropout可以有效的防止過擬合,LRN對資料進行了歸一化處理。
AlexNet網路結構如下:
input_size | out_size | kernel | stride | padding | |
---|---|---|---|---|---|
卷積層1 | [3, 227, 227] | [96, 55, 55] | (11, 11) | 4 | 0 |
池化層1 | [96, 55, 55] | [96, 27, 27] | (3, 3) | 2 | 0 |
卷積層2 | [96, 27, 27] | [256, 27, 27] | (5, 5) | 1 | 2 |
池化層2 | [256, 27, 27] | [256, 13, 13] | (3, 3) | 2 | 0 |
卷積層3 | [256, 13, 13] | [384, 13, 13] | (3, 3) | 1 | 1 |
卷積層4 | [384, 13, 13] | [384, 13, 13] | (3, 3) | 1 | 1 |
卷積層5 | [384, 13, 13] | [256, 13, 13] | (3, 3) | 1 | 1 |
池化層3 | [256, 13, 13] | [256, 6, 6] | (3, 3) | 2 | 0 |
全連線層1 | 256 * 6 * 6 | 4096 | --- | --- | --- |
input_size | out_size | kernel | stripe | padding | |
---|---|---|---|---|---|
全連線層2 | 4096 | 4096 | --- | --- | --- |
全連線層3 | 4096 | 1000 | --- | --- | --- |
計算輸出時,有一個非常重要的公式:
\[y= \frac{x-k+2p}{s} + 1 \]- y:輸出尺寸大小
- x:輸入尺寸大小
- k:卷積核大小
- p:填充數
- s:步長
構建網路模型
import torch.nn as nn
import torch.nn.functional as F
# 區域性響應歸一化
class LRN(nn.Module):
def __init__(self, local_size=1, alpha=1.0, beta=0.75, ACROSS_CHANNELS=True):
super(LRN, self).__init__()
self.ACROSS_CHANNELS = ACROSS_CHANNELS
if ACROSS_CHANNELS:
self.average=nn.AvgPool3d(kernel_size=(local_size, 1, 1),
stride=1,
padding=(int((local_size-1.0)/2), 0, 0))
else:
self.average=nn.AvgPool2d(kernel_size=local_size,
stride=1,
padding=int((local_size-1.0)/2))
self.alpha = alpha
self.beta = beta
def forward(self, x):
if self.ACROSS_CHANNELS:
div = x.pow(2).unsqueeze(1)
div = self.average(div).squeeze(1)
div = div.mul(self.alpha).add(1.0).pow(self.beta)
else:
div = x.pow(2)
div = self.average(div)
div = div.mul(self.alpha).add(1.0).pow(self.beta)
x = x.div(div)
return x
# conv
# out_size = (in_size - kernel_size + 2 * padding) / stride
class AlexNet(nn.Module):
def __init__(self):
super(AlexNet, self).__init__()
# conv
self.conv1 = nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=0)
self.conv2 = nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2)
self.conv3 = nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1)
self.conv4 = nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1)
self.conv5 = nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1)
# LRN
self.LRN = LRN(local_size=5, alpha=0.0001, beta=0.75)
# FC
self.fc1 = nn.Linear(256*6*6, 4096)
self.fc2 = nn.Linear(4096, 4096)
self.fc3 = nn.Linear(4096, 2)
# Dropout
self.Dropout = nn.Dropout()
def forward(self, x):
# conv1 -> relu -> maxpool1
# conv1: [n, 3, 227, 227] --> [n, 96, 55, 55]
# maxpool1: [n, 96, 55, 55] --> [n, 96, 27, 27]
x = F.relu(self.conv1(x))
x = self.LRN(x)
x = F.max_pool2d(x, (3, 3), 2)
# conv2 -> relu -> maxpool2
# conv2: [n, 96, 27, 27] --> [n, 256, 27, 27]
# maxpool2: [n, 256, 27, 27] --> [n, 256, 13, 13]
x = F.relu(self.conv2(x))
x = self.LRN(x)
x = F.max_pool2d(x, (3, 3), 2)
# conv3 -> relu -> conv4 -> relu
# oonv3: [n, 256, 13, 13] --> [n, 384, 13, 13]
# conv4: [n, 384, 13, 13] --> [n, 384, 13, 13]
x = F.relu(self.conv3(x))
x = F.relu(self.conv4(x))
# conv5 -> relu -> maxpool3
# conv5: [n. 384, 13, 13] --> [n, 256, 13, 13]
# maxpool3: [n, 256, 13, 13] --> [n, 256, 6, 6]
x = F.relu(self.conv5(x))
x = F.max_pool2d(x, (3, 3), 2)
# need view first for conv --> FC
x = x.view(x.size()[0], -1)
# fc1 -> fc2 -> fc3 -> softmax
# fc1: 256*6*6 --> 4096
# fc2: 4096 --> 4096
# fc3: 1000 --> 2
x = F.relu(self.fc1(x))
x = self.Dropout(x)
x = F.relu(self.fc2(x))
x = self.Dropout(x)
x = self.fc3(x)
x = F.softmax(x)
return x
由於本次實驗是一個二分類問題,所以將最後一個全連線層的輸出個數由1000改成2即可。
Transforms資料預處理
transforms定義了對資料進行怎樣的預處理,但資料的預處理並不在這裡實現,通常將transforms作為一個引數傳入自定義的資料集類,並在__ getitem __方法中實現資料的預處理。
pre_transforms = transforms.Compose([
transforms.Resize((227, 227)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
自定義資料集類
程式碼如下:
class CatDogDataset(data.Dataset):
def __init__(self, args, mode='train', transform=None):
self.args = args
self.transform = transform
self.mode = mode
self.names = self.__dataset_info()
def __getitem__(self, index):
x = imread(self.args.data_path + "/" + self.names[index], mode='RGB') # numpy
x = Image.fromarray(x) # PIL
x_label = 0 if 'cat' in self.names[index] else 1
if self.transform is not None:
x = self.transform(x)
return x, x_label
def __len__(self):
return len(self.names)
# 取train中前500張的貓和狗圖片為測試集,所以一共有1000張測試集,24000張訓練集
def __dataset_info(self):
img_path = self.args.data_path
imgs = [f for f in os.listdir(img_path) if
os.path.isfile(os.path.join(img_path, f)) and f.endswith('.jpg')]
names = []
for name in imgs:
index = int(name.split('.')[1])
# train dataset
if self.mode == 'train':
if index >= 500:
names.append(name)
# test dataset: 1000 imgs
elif self.mode == 'test':
if index < 500:
names.append(name)
return names
在類中,必須實現__ init __ ** ,__ getitem ** ,** len __** 三個方法。
在定義好了我們的資料集類之後,需要對該類進行例項化:
# get datasets
train_dataset = CatDogDataset(args, 'train', pre_transforms)
test_dataset = CatDogDataset(args, 'test', pre_transforms)
# print the length of train_dataset
print('train:{} imgs'.format(len(train_dataset)))
DataLoader
接下來要通過Pytorch自帶的DataLoader來得到一個Loader物件,該物件可以通過for ... in ...進行迭代,每一次迭代的結果就是資料集類__ getitem __ 方法返回的值。
# generate DataLoader
train_loader = DataLoader(train_dataset, args.batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, 1, shuffle=False)
使用GPU加速
個人推薦使用device的方式對Tensor進行GPU處理,因為這樣無論電腦是否安裝了CUDA+CuDNN都能不改任何程式碼成功執行,同時如果需要在另一張顯示卡上執行,只需要修改一個數字即可,很方便。
# GPU,如需指定顯示卡,只需要將0改成要指定的顯示卡的對應序號即可。
if args.without_gpu:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("use CPU !")
else:
if torch.cuda.is_available():
device = torch.device("cuda")
print("use GPU !")
else:
print("No GPU is available !")
alexnet = AlexNet()
base_epoch = 0
if args.load:
model_path = './checkpoints/99_loss_0.523277.pth'
alexnet.load_state_dict(torch.load(model_path)['alexnet'])
base_epoch = torch.load(model_path)['epoch']
# 轉換到GPU環境
alexnet.to(device)
下圖顯示了一臺伺服器上的顯示卡資訊,可以看到圖中有兩張顯示卡,其序號分別是0和1,如需使用第二張顯示卡,只需要將"cuda:0"改成"cuda:1"就可以了。
損失函式與優化方法
本次實驗,使用了交叉熵作為損失函式,隨機梯度下降法SGD作為優化方法
# loss and optim function
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(alexnet.parameters(),
lr=args.lr, momentum=0.9, weight_decay=5e-4)
迭代,對資料進行處理
首先我們需要將每次迭代生成的資訊to到相應的GPU裝置上,然後進行常規化處理:預測得到標籤,梯度清0,計算損失值,將損失值反向傳播並進行優化,程式碼如下:
for epoch in range(args.epochs):
alexnet.train()
epoch += base_epoch
epoch_loss = 0
for idx, (imgs, labels) in enumerate(train_loader):
imgs, labels = imgs.to(device), labels.to(device)
pre_labels = alexnet(imgs)
optimizer.zero_grad()
loss = criterion(pre_labels, labels.long())
loss.backward()
optimizer.step()
epoch_loss += loss.item()
print('[{}/{}][{}/{}] loss:{:.4f}'
.format(epoch+1, args.epochs, idx+1, int(len(train_dataset) / args.batch_size), loss.item()))
# save model
aver_loss = epoch_loss * args.batch_size / len(train_dataset)
state = {
'epoch': epoch,
'alexnet': alexnet.state_dict()
}
acc = eval(alexnet, test_loader, test_dataset, device)
save_model(state, './checkpoints', '{}_{:.6f}_{:.3f}.pth'.format(epoch, aver_loss, acc))
以上就是訓練一個神經網路的基本流程,下面通過一張圖對這幾部分進行整理。