【无标题】

import torch 
from torchvision.datasets import MNIST #获取测试集以及训练集数据
from torchvision import transforms
from torch import nn # 神经网络库，内置有各种函数以及模型函数
import torch.nn.functional as F #内置有各种函数，例如MSE损失函数
from torch.utils.data import  DataLoader #加载MNIST获取到的数据

对于构建线性模型类，我们的类需要继承Pytorch中的nn库中的Module类，代码如下所示：

class MNIST_Model(nn.Module):
    def __init__(self) -> None:
        super(MNIST_Model,self).__init__()
        self.fc1 = nn.Linear(28*28*1,100,bias=True) #第一层线性模型
        self.fc2 = nn.Linear(100,10,bias=True) # 第二层线性模型
    
    def forward(self,x) -> None:
        # 修改数据尺寸
        # x = x.view([-1,1*28*28])
        x = x.view([x.size(0),1*28*28])
        # 第一次全连接
        y1 = self.fc1(x)
        # 激活函数
        y1_relu = F.relu(y1)
        # 第二次全连接
        y2 = self.fc2(y1_relu)
        return y2

获取数据集需要用到另外的一个库$torchvision$库，利用该库下的函数$MNIST$数据集，

    # root即根目录，train即是否为训练集数据,download为是否下载，transformer为自定义参数。
    data_set = MNIST(root='./data/',train=True,download=True,transform=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(
        mean=(0.1307,),std=(0.3081,)
        )
    ]))
  test_loader = MNIST(root='./data/',train=False,download=True,transform=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(
        mean=(0.1307,),std=(0.3081,)
        )
    ]))

在获取到数据集之后编写测试函数以及训练函数，代码如下：

def Train_Test(epoch):
    print("-----------TRAIN-------------")
    # 获取测试集数据
    data_set = MNIST(root='./data/',train=True,download=True,transform=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(
        mean=(0.1307,),std=(0.3081,)
        )
    ]))
    # 加载数据，同时设置batch_size为1000，shuffle打乱数据
    data_loader = DataLoader(dataset=data_set,batch_size=128,shuffle=True)

    # 使用GPU进行训练
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    # 实例化线性网络
    m = MNIST_Model().to(device=device)
    # 定义优化器和损失函数
    optimizer = torch.optim.SGD(m.parameters(),lr=1e-3,momentum=1e-7)
    loss_fuc = torch.nn.CrossEntropyLoss()

    # 开始训练
    for i in range(epoch):
        for idx,(datas,labels) in enumerate(data_loader):
            # 梯度置零
            optimizer.zero_grad()
            # 前向传播
            y_out = m(datas.to(device))
            labels = labels.to(device)

            # 计算损失函数
            loss = torch.nn.functional.cross_entropy(y_out,labels)
            # 反向传播
            loss.backward()
            # 更新权重
            optimizer.step()
            if idx%1000==0:
                print("loss: ",loss.item())

    print("----------TEST------------")
    # 测试结果
    test_loss = 0
    correct = 0
    # 设置评估方法
    m.eval()
    # 获取测试集数据
    test_loader = MNIST(root='./data/',train=False,download=True,transform=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(
        mean=(0.1307,),std=(0.3081,)
        )
    ]))
    test_datas = DataLoader(shuffle=True,dataset=test_loader,batch_size=1000)
    # 设置无需梯度的模式
    with torch.no_grad():
        for data,target in test_datas:
            output = m(data.to(device))
            test_loss += F.cross_entropy(output, target.to(device), size_average=False).item()
            pred = output.data.max(1, keepdim=True)[1]
            pred = pred.to(device)
            correct += pred.eq(target.to(device).data.view_as(pred)).sum()

        test_loss /= len(test_datas.dataset)
        print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(test_datas.dataset),
            100. * correct / len(test_datas.dataset)))
        

本人自学Pytorch刚两天，如有错误请指正。