使用LeNet实现MINIST数据集的分类

1. 导入必要的包

import torch
from torch import nn
from torch.utils.data import DataLoader
import torch.nn.functional as F
from torchvision import transforms,datasets
import torch.optim as optim

2. 设置数据的transform转换并设置其他的超参数
   将数据转换为Tensor张量，然后设置Normalization

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,),(0.3081,))
])

batch_size = 256
num_workers = 2
device = ("cuda:0" if torch.cuda.is_available() else "cpu")
epochs = 10

3. 获取数据集
   如果没有数据集需要把download设置为True，这时就会自动从网上下载需要的数据集

tran_data = datasets.MNIST(root="./MNIST",
                           train=True,
                           transform=transform,
                           download=False)
test_data = datasets.MNIST(root="./MNIST",
                           train=False,
                           transform=transform,
                           download=False)

4. 加载数据集
   有了数据集需要将数据集里面的数据取出来，这时用的是DataLoader类

train_iter = DataLoader(tran_data,batch_size=batch_size,shuffle=True,num_workers=num_workers)
test_iter = DataLoader(test_data,batch_size=batch_size,shuffle=False,num_workers=num_workers)

5. 定义网络结构并设置优化器

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1,6,kernel_size=5,padding=2)
        self.conv2 = nn.Conv2d(6,16,kernel_size=5)
        self.lin1 = nn.Linear(16*5*5,120)
        self.lin2 = nn.Linear(120,84)
        self.lin3 = nn.Linear(84,10)
        self.avg = nn.AvgPool2d(kernel_size=2,stride=2)

    def forward(self,X):
        X = self.avg(F.relu(self.conv1(X)))
        X = self.avg(F.relu(self.conv2(X)))
        X = X.view(-1,400)
        X = self.lin3(self.lin2(self.lin1(X)))
        out = F.softmax(X,dim=1)
        return out

model = Net().to(device)
# 使用Adam优化器
optimizer = optim.Adam(model.parameters())

6. 进行模型的训练

def train_model(epoch):
    model.train()
    for batch_idx,(X,y) in enumerate(train_iter):
        X,y = X.to(device),y.to(device)
        y_hat = model(X)
        loss = F.cross_entropy(y_hat,y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if batch_idx % 3000 == 0:
            print("Train Epoch : {} \t Loss : {:.6f}".format(epoch, loss.item()))

def test_model():
    model.eval()

    correct = 0.0
    test_loss = 0.0

    with torch.no_grad():
        for X,y in test_iter:
            X,y = X.to(device),y.to(device)
            y_hat = model(X)
            test_loss += F.cross_entropy(y_hat,y).item()
            pred = y_hat.argmax(dim=1)
            correct += pred.eq(y.view_as(pred)).sum().item()
        test_loss /= len(test_iter.dataset)
        print("Test —— Average loss : {:.4f}, Accuracy : {:.3f}\n".format(test_loss,
                                                                      100.0 * correct / len(test_iter.dataset)))

def main():
    for epoch in range(epochs):
        train_model(epoch)
        test_model()

if __name__ == '__main__':
    main()

在实现的过程中发现了几个易错的bug

1. 首先如果需要使用CPU的多线程，也就是设置num_workers>1，那么这个时候就需要使用main函数，然后再调用main函数，否则会报错
2. 在定义网络类的时候如果没有使用super（）或者init这个函数拼写错误那么就会报Adam优化器并没有接受到参数的错误。