AlexNet训练CIFA10数据集——pytorch实现

 代码在kaggle上跑，精度最终为83%

输入一张3*227*227的图片，每一层的输出：

Conv2d output shape:	 torch.Size([1, 96, 57, 57])
ReLU output shape:	 torch.Size([1, 96, 57, 57])
MaxPool2d output shape:	 torch.Size([1, 96, 28, 28])
Conv2d output shape:	 torch.Size([1, 256, 28, 28])
ReLU output shape:	 torch.Size([1, 256, 28, 28])
MaxPool2d output shape:	 torch.Size([1, 256, 13, 13])
Conv2d output shape:	 torch.Size([1, 384, 13, 13])
ReLU output shape:	 torch.Size([1, 384, 13, 13])
Conv2d output shape:	 torch.Size([1, 384, 13, 13])
ReLU output shape:	 torch.Size([1, 384, 13, 13])
Conv2d output shape:	 torch.Size([1, 256, 13, 13])
ReLU output shape:	 torch.Size([1, 256, 13, 13])
MaxPool2d output shape:	 torch.Size([1, 256, 6, 6])
Flatten output shape:	 torch.Size([1, 9216])
Linear output shape:	 torch.Size([1, 4096])
ReLU output shape:	 torch.Size([1, 4096])
Dropout output shape:	 torch.Size([1, 4096])
Linear output shape:	 torch.Size([1, 4096])
ReLU output shape:	 torch.Size([1, 4096])
Dropout output shape:	 torch.Size([1, 4096])
Linear output shape:	 torch.Size([1, 10])

import os
import datetime
import torch
import torchvision
from torch import nn
from torch import optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
import torchvision.models as models
from torchvision.utils import save_image
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
from PIL import Image
import time
import argparse

def try_gpu(i=0):
    """如果存在，则返回gpu(i)，否则返回cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

batch_size = 128
path = './'
train_transform = transforms.Compose([
    transforms.RandomSizedCrop(227),# 随机剪切成227*227
    transforms.RandomHorizontalFlip(),# 随机水平翻转
    transforms.ToTensor(),
    transforms.Normalize(mean = [ 0.485, 0.456, 0.406 ],
                         std = [ 0.229, 0.224, 0.225 ]),
])
val_transform = transforms.Compose([
    transforms.Resize((227,227)),
    transforms.ToTensor(),
    transforms.Normalize(mean = [ 0.485, 0.456, 0.406 ],
                         std = [ 0.229, 0.224, 0.225 ]),
])

traindir = os.path.join(path, 'train')
valdir = os.path.join(path, 'val')

train_set = torchvision.datasets.CIFAR10(
    traindir, train=True, transform=train_transform, download=True)
valid_set = torchvision.datasets.CIFAR10(
    valdir, train=False, transform=val_transform, download=True)


train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
valid_loader = DataLoader(valid_set, batch_size=batch_size, shuffle=False)

dataloaders = {
    'train': train_loader,
    'valid': valid_loader,
#     'test': dataloader_test
    }

dataset_sizes = {
    'train': len(train_set),
    'valid': len(valid_set),
#     'test': len(test_set)
    }
print(dataset_sizes)

class AlexNet(nn.Module):
 
    def __init__(self):
        super().__init__()
        
        self.conv = nn.Sequential(
                    nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=4), nn.ReLU(),
                    nn.MaxPool2d(kernel_size=3, stride=2),
                    nn.Conv2d(96, 256, kernel_size=5, padding=2), nn.ReLU(),
                    nn.MaxPool2d(kernel_size=3, stride=2),
                    nn.Conv2d(256, 384, kernel_size=3, padding=1), nn.ReLU(),
                    nn.Conv2d(384, 384, kernel_size=3, padding=1), nn.ReLU(),
                    nn.Conv2d(384, 256, kernel_size=3, padding=1), nn.ReLU(),
                    nn.MaxPool2d(kernel_size=3, stride=2))
        
        self.fc = nn.Sequential(
                        nn.Linear(9216, 4096), nn.ReLU(),
                        nn.Dropout(p=0.5),
                        nn.Linear(4096, 4096), nn.ReLU(),
                        nn.Dropout(p=0.5),
                        nn.Linear(4096, 10))
        
        self.fn = nn.Flatten()

    def forward(self, x):
        out = self.conv(x)
        out = self.fn(out)
        out = self.fc(out)
        return out

# 测试
net = nn.Sequential(
    # 这里，我们使用一个11*11的更大窗口来捕捉对象。


    # 同时，步幅为4，以减少输出的高度和宽度。


    # 另外，输出通道的数目远大于LeNet
    nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=4), nn.ReLU(),
#     nn.BatchNorm2d(96),
    nn.MaxPool2d(kernel_size=3, stride=2),
    # 减小卷积窗口，使用填充为2来使得输入与输出的高和宽一致，且增大输出通道数
    nn.Conv2d(96, 256, kernel_size=5, padding=2), nn.ReLU(),
#     nn.BatchNorm2d(256),
    nn.MaxPool2d(kernel_size=3, stride=2),
    # 使用三个连续的卷积层和较小的卷积窗口。


    # 除了最后的卷积层，输出通道的数量进一步增加。


    # 在前两个卷积层之后，汇聚层不用于减少输入的高度和宽度
    nn.Conv2d(256, 384, kernel_size=3, padding=1), nn.ReLU(),
#     nn.BatchNorm2d(384),
    nn.Conv2d(384, 384, kernel_size=3, padding=1), nn.ReLU(),
#     nn.BatchNorm2d(384),
    nn.Conv2d(384, 256, kernel_size=3, padding=1), nn.ReLU(),
#     nn.BatchNorm2d(256),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Flatten(),
    # 这里，全连接层的输出数量是LeNet中的好几倍。

使用dropout层来减轻过拟合
    nn.Linear(9216, 4096), nn.ReLU(),
#     nn.BatchNorm1d(4096),
    nn.Dropout(p=0.5),
    nn.Linear(4096, 4096), nn.ReLU(),
#     nn.BatchNorm1d(4096),
    nn.Dropout(p=0.5),
    # 最后是输出层。

由于这里使用Fashion-MNIST，所以用类别数为10，而非论文中的1000
    nn.Linear(4096, 10)
#     ,nn.Softmax(dim=1)
)

X = torch.randn(1, 3, 227, 227)
for layer in net:
    X=layer(X)
    print(layer.__class__.__name__,'output shape:\t',X.shape)

def train(model, criterion, optimizer, scheduler, device, num_epochs, dataloaders,dataset_sizes):
    model = model.to(device)
    print('training on ', device)
    since = time.time()

    best_model_wts = []
    best_acc = 0.0

    for epoch in range(num_epochs):
        # 训练模型
        s = time.time()
        model,train_epoch_acc,train_epoch_loss = train_model(
            model, criterion, optimizer, dataloaders['train'], dataset_sizes['train'], device)
        print('Epoch {}/{} - train Loss: {:.4f}  Acc: {:.4f}  Time:{:.1f}s'
            .format(epoch+1, num_epochs, train_epoch_loss, train_epoch_acc,time.time()-s))
        # 验证模型
        s = time.time()
        val_epoch_acc,val_epoch_loss = val_model(
            model, criterion, dataloaders['valid'], dataset_sizes['valid'], device)
        print('Epoch {}/{} - valid Loss: {:.4f}  Acc: {:.4f}  Time:{:.1f}s'
            .format(epoch+1, num_epochs, val_epoch_loss, val_epoch_acc,time.time()-s))
        # 每轮都记录最好的参数.
        if val_epoch_acc > best_acc:
            best_acc = val_epoch_acc
            best_model_wts = model.state_dict()
        # 优化器
#         if scheduler not in None:
#             scheduler.step()
        # 保存画图参数
        train_losses.append(train_epoch_loss.to('cpu'))
        train_acc.append(train_epoch_acc.to('cpu'))
        val_losses.append(val_epoch_loss.to('cpu'))
        val_acc.append(val_epoch_acc.to('cpu'))
        print()
    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))
    # model.load_state_dict(best_model_wts)
    return model

def train_model(model, criterion, optimizer, dataloader, dataset_size,device):
    model.train()
    running_loss = 0.0
    running_corrects = 0
    for inputs,labels in dataloader:
        optimizer.zero_grad()
        # 输入的属性
        inputs = Variable(inputs.to(device))
        # 标签
        labels = Variable(labels.to(device))
        # 预测
        outputs = model(inputs)
        _,preds = torch.max(outputs.data,1)
        # 计算损失
        loss = criterion(outputs,labels)
        #梯度下降
        loss.backward()
        optimizer.step()

        running_loss += loss.data
        running_corrects += torch.sum(preds == labels.data)

    epoch_loss = running_loss / dataset_size
    epoch_acc = running_corrects / dataset_size
    
    return model,epoch_acc,epoch_loss

def val_model(model, criterion, dataloader, dataset_size, device):
    model.eval()
    running_loss = 0.0
    running_corrects = 0
    for (inputs,labels) in dataloader:
        # 输入的属性
        inputs = Variable(inputs.to(device))
        # 标签
        labels = Variable(labels.to(device))
        # 预测
        outputs = model(inputs)
        _,preds = torch.max(outputs.data,1)
        # 计算损失
        loss = criterion(outputs,labels)
        
        running_loss += loss.data
        running_corrects += torch.sum(preds == labels.data)

    epoch_loss = running_loss / dataset_size
    epoch_acc = running_corrects / dataset_size
    
    return epoch_acc,epoch_loss

val_losses,val_acc = [],[]
train_losses,train_acc = [],[]

lr,num_epochs = 0.05,20
model = AlexNet()
criterion = nn.CrossEntropyLoss()
# optimizer = optim.Adam(model.parameters(), lr=lr)
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
# scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

model = train(model, criterion, optimizer, None ,
              try_gpu(), num_epochs, dataloaders, dataset_sizes)

plt.plot(range(1, len(train_losses)+1),train_losses, 'b', label='training loss')
plt.plot(range(1, len(val_losses)+1), val_losses, 'r', label='val loss')
plt.legend()

plt.plot(range(1,len(train_acc)+1),train_acc,'b--',label = 'train accuracy')
plt.plot(range(1,len(val_acc)+1),val_acc,'r--',label = 'val accuracy')
plt.legend()