使用cifar10数据集和alexnet网络模型训练分类模型

下载cifar10数据集

代码：

import torchvision
import torch
transform = torchvision.transforms.Compose(
    [torchvision.transforms.ToTensor(),
     torchvision.transforms.Resize(224)]
)

train_set = torchvision.datasets.CIFAR10(root='./',download=False,train=True,transform=transform)
test_set = torchvision.datasets.CIFAR10(root='./',download=False,train=False,transform=transform)
train_loader = torch.utils.data.DataLoader(train_set,batch_size=8,shuffle=True)
test_loader = torch.utils.data.DataLoader(test_set,batch_size=8,shuffle=True)

class Alexnet(torch.nn.Module):  #1080 2080
    def __init__(self,num_classes=10):
        super(Alexnet,self).__init__()
        net = torchvision.models.alexnet(pretrained=False)  #迁移学习
        net.classifier = torch.nn.Sequential()
        self.features = net
        self.classifier = torch.nn.Sequential(
            torch.nn.Dropout(0.3),
            torch.nn.Linear(256 * 6 * 6, 4096),
            torch.nn.ReLU(inplace=True),
            torch.nn.Dropout(0.3),
            torch.nn.Linear(4096, 4096),
            torch.nn.ReLU(inplace=True),
            torch.nn.Linear(4096, num_classes),
        )
    def forward(self,x):
        x = self.features(x)
        x = x.view(x.size(0),-1)
        x = self.classifier(x)

        return x
device = torch.device('cpu')
net = Alexnet().to(device)
loss_func = torch.nn.CrossEntropyLoss().to(device)
optim = torch.optim.Adam(net.parameters(),lr=0.001)

net.train()
for epoch in range(10):
    for step,(x,y) in enumerate(train_loader):  # 28*28*1  32*32*3
        x,y = x.to(device),y.to(device)
        output = net(x)
        loss = loss_func(output,y)
        optim.zero_grad()
        loss.backward()
        optim.step()
    print("epoch:",epoch,'loss:',loss)