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)