【pytorch图像分类】mnist手写数据集识别

对比多方教程，自己总结如何入门使用pytorch学习搭建基础网络模型进行训练和测试

目录

1.准备数据

2.超参数的设定

3.创建网络 

4.优化器的定义

5.训练定义

6.测试定义

7.测试模型

总结：

1.准备数据

这部分将会用到相关dataset和Dataloader2

#.1准备数据，这些需要准备dataset.DataLoader
  #数据的准备
train_data =datasets.MNIST("./data/train",train=True,transform= torchvision.transforms.ToTensor(),
                                       download=True)

test_data = datasets.MNIST("./data/test",train=False,transform= torchvision.transforms.ToTensor()
                            ,download=True)

print("训练集的长度：{}".format(len(train_data)))
print("测试集的长度：{}".format(len(test_data)))

  #数据迭代器的准备
train_loader=  DataLoader(train_data,batch_size=4,shuffle=True)
test_loader = DataLoader(test_data,batch_size=4)

2.超参数的设定

#2.超参数的设定
epoch = 20
bactch_size =16
device = "cuda" if torch.cuda.is_available() else "cpu"
start_time = time.time()

3.创建网络 

#3.构建模型，这里可以使用torch构造一个深层的神经网络
class MnistNet(nn.Module):
    def __init__(self):
        super(MnistNet, self).__init__()
        # self.fc1 = nn.Linear(28 * 28 * 1, 28)  # 定义Linear的输入和输出的形状
        # self.fc2 = nn.Linear(28, 10)  # 定义Linear的输入和输出的形状

        self.flatten = nn.Flatten()  #从第一维到最后一维
        self.liner_relu_nn = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

    def  forward(self,x):
        # x = x.view(-1,28*28*1)  #拉直  展平相当于flatten
        x = self.flatten(x)
        Y = self.liner_relu_nn(x)
        return  Y



 #实例化模型
modelNet = MnistNet().to(device)
# print(model)

4.优化器的定义

#4.优化器的定义
loss_fn = nn.CrossEntropyLoss().to(device)   #交叉熵损失函数
optimizer = torch.optim.SGD(modelNet.parameters() , lr = 1e-4)  #SGD优化器

 5.训练定义

#5. 模型训练的定义
def train_mnist(train_loader,model,loss_fn,optimizer):
    '''训练过程写成函数'''
    modelNet.train(True)  #nn.module.train()方法表示训练中会使用到Dropout，BN等方法

    total_train_step = 0  #记录训练轮次
    #idx次数，也就是img的索引
    for data in train_loader:
        img , target = data
        img,target = img.to(device),target.to(device)

        pred = model(img)    #将图片输入网络预测

        train_loss = loss_fn(pred,target)  #损失函数

        #反向传播
        optimizer.zero_grad()  #梯度清零
        train_loss.backward()
        optimizer.step()  #参数的优化更新

        total_train_step += 1
        if total_train_step %1000 == 0:
            print("训练的轮次：{}，loss：{}".format(total_train_step,train_loss.item()))

6.测试定义

#6.模型的验证
def test_mnist(test_loader,model,device):
    #模型的验证
    model.eval()  #eval不会使用dropout，BN等方法，符合测试的要求
    correct = 0   #总和的正确率
    test_loss = 0 #test总的测试损失
    with torch.no_grad():   #不用也不会计算梯度，也不会进行反向传播
        for data , target in test_loader:
            data,target = data.to(device) , target.to(device)
            #得到测试结果
            pred = model(data)
            #计算测试损失
            test_loss  += loss_fn(pred,target).item()  #必须加item

            #找到概率最大的下标
            # max_pred = torch.max(pred,dim=1)
            # max_pred = pred.argmax(dim=1)
            #累计正确的值
            correct =correct +  (pred.argmax(dim=1) == target)   #土堆方法
            # correct += max_pred.eq(target.view_as(max_pred)).sum().item() #官方方法
            #整个测试集的正确率
        test_loss = test_loss /len(test_loader)
        test_correct =  correct / len(test_loader)
        print("Test---Average:{}   Accuracy:{}".format(test_loss,test_correct))

 7.测试模型

#7.调用train_mnist和test_mnist方法
for i in range(epoch):
    print("----------第{}轮训练开始-------".format(i + 1))
    train_mnist(train_loader, modelNet, loss_fn, optimizer)
    test_mnist(test_loader,modelNet,device)


#5.模型的保存
# torch.save(MnistNet,"Mnist.pth")