本文已参与「新人创作礼」活动， 一起开启掘金创作之路。

使用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)

Lenet5训练模型

下载数据集

可以提前下载也可以在线下载

train_data = torchvision.datasets.MNIST(root='./',download=True,train=True,transform=transform)
test_data = torchvision.datasets.MNIST(root='./',download=True,train=False,transform=transform)

训练模型

import torch
import torchvision
class Lenet5(torch.nn.Module):
    def __init__(self):
        super(Lenet5, self).__init__()
        self.model = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=1,out_channels=6,kernel_size=5), # 1*32*32 # 6*28*28
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2,  # 6*14*14
                                stride=2),
            torch.nn.Conv2d(in_channels=6,
                                    out_channels=16,
                                    kernel_size=5),  # 16 *10*10
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2,
                                stride=2),  # 16*5*5
            torch.nn.Conv2d(in_channels=16,
                                    out_channels=120,
                                    kernel_size=5),  # 120*1*1
            torch.nn.ReLU(),
            torch.nn.Flatten(),  # 展平 成120*1维
            torch.nn.Linear(120, 84),
            torch.nn.Linear(84, 10)
        )

    def forward(self,x):
        x = self.model(x)
        return x

transform = torchvision.transforms.Compose(
    [torchvision.transforms.Resize(32),
     torchvision.transforms.ToTensor()]
)
train_data = torchvision.datasets.MNIST(root='./',download=True,train=True,transform=transform)
test_data = torchvision.datasets.MNIST(root='./',download=True,train=False,transform=transform)

#分批次加载数据  64 128
train_loader =torch.utils.data.DataLoader(train_data,batch_size=64,shuffle=True)
test_loader =torch.utils.data.DataLoader(test_data,batch_size=64,shuffle=True)


#gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = Lenet5().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):
        x = x.to(device)
        y = y.to(device)

        ouput = net(x)
        loss = loss_func(ouput,y)  #计算损失

        optim.zero_grad()
        loss.backward()
        optim.step()

    print('epoch:',epoch,"loss:",loss)

torch.save(net,'net.pkl')

import torch
import torchvision
class Lenet5(torch.nn.Module):
    def __init__(self):
        super(Lenet5, self).__init__()
        self.model = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=1,out_channels=6,kernel_size=5), # 1*32*32 # 6*28*28
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2,  # 6*14*14
                                stride=2),
            torch.nn.Conv2d(in_channels=6,
                                    out_channels=16,
                                    kernel_size=5),  # 16 *10*10
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2,
                                stride=2),  # 16*5*5
            torch.nn.Conv2d(in_channels=16,
                                    out_channels=120,
                                    kernel_size=5),  # 120*1*1
            torch.nn.ReLU(),
            torch.nn.Flatten(),  # 展平 成120*1维
            torch.nn.Linear(120, 84),
            torch.nn.Linear(84, 10)
        )

    def forward(self,x):
        x = self.model(x)
        return x


transform = torchvision.transforms.Compose(
    [torchvision.transforms.Resize(32),
     torchvision.transforms.ToTensor()]
)
test_data = torchvision.datasets.MNIST(root='./',download=False,train=False,transform=transform)
test_loader =torch.utils.data.DataLoader(test_data,batch_size=64,shuffle=True)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = torch.load('net.pkl')

net.eval()  #表明进行推理
with torch.no_grad():
    for step,(x,y) in enumerate(test_loader):
        x,y = x.to(device),y.to(device)

        pre = net(x)
        print(pre)
        pre_y = torch.max(pre.cpu(),1)[1].numpy()
        print(pre_y)
        y = y.cpu().numpy()

        acc = (pre_y == y).sum()/len(y)

    print("accu:",acc)