[pytorch] 利用Alexnet训练cifar10

原创

已于 2022-09-05 11:40:44 修改 · 7.7k 阅读

65 ·

CC 4.0 BY-SA版权

文章标签：

#python #pytorch

于 2022-03-30 19:29:05 首次发布

本文介绍如何使用PyTorch框架在CIFAR10数据集上训练AlexNet模型，并展示了详细的训练过程及结果。通过调整图片尺寸和输出类别数量等参数，实现了模型的有效训练。

论文地址：ImageNet Classification with Deep Convolutional Neural Networks
Alexnet的网络结构如图所示
Alexnet结构
cifar10数据集一共包含10个类别的RGB彩色图像，每个类别分别有6000张，共60000张大小为32x32的图片。其中50000个作为训练集，10000个作为测试集。

代码如下：
注意这里将cifar10图片上采样为224*224，以接近原始版本Alexnet；输出类别由1000修改为10。

# 导入模块
import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
from torchvision.transforms.functional import InterpolationMode
import random
import matplotlib.pyplot as plt

# 下载以及转换cifar10
transform = transforms.Compose([transforms.Resize((224, 224), interpolation=InterpolationMode.BICUBIC),
                                transforms.ToTensor(),
                                transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.247, 0.2435, 0.2616]) # 此为训练集上的均值与方差
                                ])
train_images = datasets.CIFAR10('./', train=True, download=True, transform=transform)
test_images = datasets.CIFAR10('./', train=False, download=True, transform=transform)

# batch size设置为256
train_data = DataLoader(train_images, batch_size=256, shuffle=True, num_workers=2)
test_data = DataLoader(test_images, batch_size=256, num_workers=2)

# Alexnet
class Model(nn.Module):
  def __init__(self):
    super().__init__()
    self.net = nn.Sequential(nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=1), nn.ReLU(),
                             nn.MaxPool2d(kernel_size=3, stride=2), 
                             nn.Conv2d(96, 256, kernel_size=5,