(Pytorch)AlexNet代码复现CIFAR-10数据集
model.py
import torch.nn as nn
import torch
from torch.utils.tensorboard import SummaryWriter
num_classes = 10
class AlexNet(nn.Module):
def __init__(self):
super(AlexNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2), # input[3, 224, 224] output[48, 55, 55]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[48, 27, 27]
nn.Conv2d(48, 128, kernel_size=5, padding=2), # output[128, 27, 27]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[128, 13, 13]
nn.Conv2d(128, 192, kernel_size=3, padding=1), # output[192, 13, 13]
nn.ReLU(inplace=True),
nn.Conv2d(192, 192, kernel_size=3, padding=1), # output[192, 13, 13]
nn.ReLU(inplace=True),
nn.Conv2d(192, 128, kernel_size=3, padding=1), # output[128, 13, 13]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[128, 6, 6]
)
self.classifier = nn.Sequential(
nn.Dropout(p=0.5), # Dropout 随机失活神经元，默认比例为0.5 (AlexNet防止过拟合的手段)
nn.Linear(128 * 6 * 6, 2048),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(2048, 2048),
nn.ReLU(inplace=True),
nn.Linear(2048, num_classes),
)
# 前向传播
def forward(self, x):
x = self.features(x)
x = torch.flatten(x, start_dim=1) # 展平后再传入全连接层
x = self.classifier(x)
return x
# writer = SummaryWriter("model_logs")
if __name__ == '__main__':
model = AlexNet()
print(model)
# writer.add_graph(model, torch.rand(1, 3, 224, 224))
# writer.close()train.py
import os
import sys
import torch
import torchvision
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import transforms
from model import AlexNet
from tqdm import tqdm
from torch.utils.tensorboard import SummaryWriter
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("using {} device.".format(device))
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
"val": transforms.Compose([transforms.Resize((224, 224)), # cannot 224, must (224, 224)
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}
train_dataset = torchvision.datasets.CIFAR10(root='data', train=True, download=True, transform=data_transform["train"])
val_dataset = torchvision.datasets.CIFAR10(root='data', train=False, transform=data_transform["val"])
# {'airplane': 0, 'automobile': 1, 'bird': 2, 'cat': 3, 'deer': 4, 'dog': 5, 'frog': 6, 'horse': 7, 'ship': 8,
batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False)
train_num = len(train_dataset)
val_num = len(val_dataset)
model = AlexNet()
model.to(device)
loss_function = nn.CrossEntropyLoss()
loss_function = loss_function.to(device)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
best_acc = 0.0
epochs = 200
save_path = 'AlexNet.pth'
train_steps = len(train_loader)
val_steps = len(train_loader)
writer = SummaryWriter("logs")
resume = True # 设置是否需要从上次的状态继续训练
if resume:
if os.path.isfile("AlexNet.pth"):
print("Resume from checkpoint...")
checkpoint = torch.load("AlexNet.pth")
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
initepoch = checkpoint['epoch'] + 2
print("====>loaded checkpoint (epoch{})".format(checkpoint['epoch']+1))
else:
print("====>no checkpoint found.")
initepoch = 1 # 如果没进行训练过，初始训练epoch值为1
for epoch in range(initepoch - 1, epochs):
# train
print("-------第 {} 轮训练开始-------".format(epoch + 1))
model.train()
train_acc = 0.0
running_loss = 0.0
train_bar = tqdm(train_loader, file=sys.stdout)
for step, data in enumerate(train_bar):
images, labels = data
optimizer.zero_grad()
outputs = model(images.to(device))
loss = loss_function(outputs, labels.to(device))
loss.backward()
optimizer.step()
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1, epochs, loss)
_, predict = torch.max(outputs, dim=1)
train_acc += torch.eq(predict, labels.to(device)).sum().item()
train_loss = running_loss / train_steps
train_accurate = train_acc / train_num
# val
model.eval()
val_acc = 0.0
running_loss = 0.0
with torch.no_grad():
val_bar = tqdm(val_loader, file=sys.stdout)
for step, val_data in enumerate(val_bar):
val_images, val_labels = val_data
outputs = model(val_images.to(device))
loss = loss_function(outputs, val_labels.to(device))
running_loss += loss.item()
_, predict = torch.max(outputs, dim=1)
val_acc += torch.eq(predict, val_labels.to(device)).sum().item()
val_loss = running_loss / val_steps
val_accurate = val_acc / val_num
print('[epoch %d] train_loss: %.3f val_loss:%.3f train_accuracy:%.3f val_accuracy: %.3f' %
(epoch + 1, train_loss,val_loss,train_accurate, val_accurate))
writer.add_scalars('loss',
{'train': train_loss, 'val': val_loss}, global_step=epoch)
writer.add_scalars('acc',
{'train': train_accurate, 'val': val_accurate}, global_step=epoch)
# 保存断点
checkpoint = {"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"epoch": epoch}
path_checkpoint = "AlexNet.pth"
torch.save(checkpoint, path_checkpoint)
print("保存模型成功")
print('Finished Training')
writer.close()
程序设置了断点续训，停止后可继续训练，日志文件可以用tensorboard查看。