pytorch手写数字识别的案例

整个代码可以分为数据、网络定义、loss、optimizer、训练和测试、模型保存这几个部分。

数据读取和加载

import torch
import torchvision.datasets as dataset
#数据
#读取mnist数据
train_data = dataset.MNIST(root='mnist',
                           train=True,
                           transform=transforms.ToTensor(),
                           download=True)
test_data = dataset.MNIST(root='mnist',
                          train=False,
                          transform=transforms.ToTensor(),
                          download=False)

#定义数据加载，每次一个batch
train_loader = data_utils.DataLoader(dataset=train_data,
                                     batch_size=64,
                                     shuffle=True)

test_loader = data_utils.DataLoader(dataset=test_data,
                                     batch_size=64,
                                     shuffle=True)

网络定义和初始化：

class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN,self).__init__()
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(1,32,kernel_size=5,padding=2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2))
        self.fc = torch.nn.Linear(14*14*32,10)

    def forward(self,x):
        x = self.conv(x)
        #输出要转换为batch X -1 维的向量才好传入fc层
        t = x.view(x.size()[0],-1)
        out = self.fc(t)
        return out
#初始化
cnn = CNN()
cnn = cnn.cuda()

loss

#loss,这里使用交叉熵
loss_func = torch.nn.CrossEntropyLoss()

optimizer

#optimizer
optimizer = torch.optim.Adam(cnn.parameters(),lr=0.01)

训练和测试

for epoch in range(10):
    for i, (images,labels) in enumerate(train_loader):
        #数据传到cuda上
        images = images.cuda()
        labels = labels.cuda()

        outputs = cnn(images)

        loss = loss_func(outputs,labels)

        optimizer.zero_grad()

        loss.backward()
        optimizer.step()

        print("epoch is {},ite is {}/{}.loss is {}".format(epoch+1,i,len(train_data)//64,loss.item()))

        #test
        loss_test=0
        accuracy=0
        for i, (images,labels) in enumerate(test_loader):
            images = images.cuda()
            labels = labels.cuda()
            outputs = cnn(images)
            loss_test += loss_func(outputs,labels)
            #outputs中最大的一个值就是分类结果
            _,pred = outputs.max(1)
            accuracy += (pred == labels).sum().item()

        accuracy = accuracy/len(test_data)
        #这种写法比较笨，如果bachsize修改了这儿忘记改了就会报错。过两天我再写使用argparse管理参数的代码
        loss_test = loss_test/(len(test_data)//64)

        print("epoch is {},accuracy is {},loss test is {}".format(epoch+1,accuracy,loss_test.item()))

保存

torch.save(cnn.state_dict(),"model/mnist_cls.pkl")