深度学习训练测试全过程

完整代码:

import torch.cuda
from torchvision import datasets
from torchvision.transforms import ToTensor


'''下载训练集MNIST(图片+标签)'''
training_data = datasets.MNIST(
    root="data", # 数据集下载的根目录
    train=True,  # 下载的数据集是否是训练集，为True数据集赋予training_data变量
    download=True, # 是否下载
    transform=ToTensor(), # 训练集转换为张量
)

'''下载训练集MNIST(图片+标签)'''
test_data = datasets.MNIST(
    root="data", # 数据集下载的根目录
    train=False,  # 下载的数据集是否是测试集，为False数据集赋予test_data变量
    download=True, # 是否下载
    transform=ToTensor(), # 训练集转换为张量
)

from torch.utils.data import DataLoader
# 加载数据集
train_dataloader = DataLoader(training_data, batch_size=64) # 每次加载64图片为一批次
test_dataloader = DataLoader(test_data, batch_size=64) # 每次加载64图片为一批次

from torch import nn
'''定义神经网络'''
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        # Flatten 数据扁平化
        self.flatten = nn.Flatten()
        self.hidden1 = nn.Linear(28*28,128)
        self.hidden2 = nn.Linear(128,256)
        self.out = nn.Linear(256,10)
    def forward(self,x):
        x = self.flatten(x)
        x = self.hidden1(x)
        x = torch.relu(x)
        x = self.hidden2(x)
        x = torch.relu(x)
        x = self.out(x)
        return x


# 判断是否存在cuda
device = "cuda" if torch.cuda.is_available() else "cpu"
# 模型实例化，并在device设备上运行
model = NeuralNetwork().to(device)
# 输出模型
print(model)


# 定义损失函数
loss_fn = nn.CrossEntropyLoss() # 交叉熵损失函数,可以进行多分类
# 定义优化器
optimizer = torch.optim.Adam(model.parameters(),lr=0.001)
def train(dataloader,model,loss_fn,optimizer):
    # 切换模式为训练，另一个模式是model.eval()
    model.train()
    # 设置批次序数
    batch_size_num = 1
    # 对训练集的图片和标签进行遍历训练
    for imgs,labels in dataloader:
        # 数据传入到设备(CPU或GPU)
        imgs,labels = imgs.to(device),labels.to(device)
        # 前向传播
        pred = model.forward(imgs)
        # 损失值计算
        loss = loss_fn(pred,labels)
        # 梯度清零,之前的梯度清除
        optimizer.zero_grad()
        # 根据损失值求梯度，反向传播，得到每个参数的梯度值
        loss.backward()
        # 根据梯度更新参数
        optimizer.step()

        # 获取上述的损失值并打印
        loss = loss.item()
        # print(f"loss:{loss:>7f} [number:{batch_size_num}]")
        batch_size_num+=1

def test(dataloader,model,loss_fn):
    # 获取整个数据集的长度
    size = len(dataloader.dataset)
    # 获取整个的批次数量
    num_batches = len(dataloader)
    # 切换模式为测试模式
    model.eval()
    # 损失值和正确率初始化
    test_loss,correct = 0,0
    # 关闭梯度,减少内存开销,当前用不到backward反向传播的时候
    with torch.no_grad():
        # 对训练集的图片和标签进行遍历训练
        for imgs, labels in dataloader:
            # 数据传入到设备(CPU或GPU)
            imgs, labels = imgs.to(device), labels.to(device)
            # 前向传播
            pred = model.forward(imgs)
            # 测试损失值
            test_loss += loss_fn(pred,labels).item()
            # 对所有图片提取预测中概率最大的那个数值的下标是否等于真实值的标签,argmax()参数1代表行,0代表列
            # 再将bool值转换为浮点型数据，累加为1的值代表正确预测值之和
            correct += (pred.argmax(1) == labels).type(torch.float).sum().item()
    # 批次测试损失值等于所有批次的损失值除以批次数
    test_loss /=num_batches
    # 正确率，所有图片预测正确值除以总图片的个数
    correct /=size
    print(f"测试结果: \n 精确度:{(100*correct)}%,平均损失: {test_loss}")