pyTorch入门（一）——Minist手写数据识别训练全连接网络

import torch
import torch.nn.functional as F

##Minist的图像为1X28X28的
class LinearNet(torch.nn.Module):
    def __init__(self):
        super(LinearNet, self).__init__()
        ##784是图像为 1X28X28，通道X宽度X高度，然后总的输入就是784
        self.l1 = torch.nn.Linear(784, 512)
        self.l2 = torch.nn.Linear(512, 256)
        self.l3 = torch.nn.Linear(256, 128)
        self.l4 = torch.nn.Linear(128, 64)
        self.l5 = torch.nn.Linear(64, 10)
        ##定义损失函数
        self.criterion = torch.nn.CrossEntropyLoss()


    def forward(self, x):
        ##将输入的图像矩阵改为N行784列
        x = x.view(-1, 784)
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        x = F.relu(self.l3(x))
        x = F.relu(self.l4(x))
        ##最后一层激活在损失函数中加入了，这里直接输出，不要加上rule了
        return self.l5(x)

import torch
import time
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.optim as optim
from NetLinear import LinearNet

batch_size = 64
##设置本次要训练用的模型
train_name = 'LinearNet'
print("train_name:" + train_name)
##设置模型保存名称
savemodel_name = train_name + ".pt"
print("savemodel_name:" + savemodel_name)
##设置初始预测率,用于判断高于当前预测率的保存模型
toppredicted = 0.0
##设置学习率
learnrate = 0.01 
##设置动量值，如果上一次的momentnum与本次梯度方向是相同的，梯度下降幅度会拉大，起到加速迭代的作用
momentnum = 0.5

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.1307,), std=(0.3081,))
]) ##Normalize 里面两个值0.1307是均值mean， 0.3081是标准差std，计算好的直接用了

##训练数据集位置，如果不存在直接下载
train_dataset = datasets.MNIST(
    root = '../datasets/mnist', 
    train = True,
    download = True,
    transform = transform
)
##读取训练数据集
train_dataloader = DataLoader(
    dataset= train_dataset,
    shuffle=True,
    batch_size=batch_size
)
##测试数据集位置，如果不存在直接下载
test_dataset = datasets.MNIST(
    root= '../datasets/mnist',
    train= False,
    download=True,
    transform= transform
)
##读取测试数据集
test_dataloader = DataLoader(
    dataset= test_dataset,
    shuffle= True,
    batch_size=batch_size
)

##设置选择训练模型,因为python用的是3.9,用不了match case语法
def switch(train_name):
    if train_name == 'LinearNet':
        return LinearNet()


##定义训练模型
class Net(torch.nn.Module):
    def __init__(self, train_name):
        super(Net, self).__init__()
        self.model = switch(train_name= train_name)
        self.criterion = self.model.criterion

    def forward(self, x):
        x = self.model(x)
        return x


model = Net(train_name)
##加入判断是CPU训练还是GPU训练
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

##优化器 
optimizer = optim.SGD(model.parameters(), lr= learnrate, momentum= momentnum)

##训练函数
def train(epoch):
    running_loss = 0.0
    for batch_idx, data in enumerate(train_dataloader, 0):
        inputs, target = data
        ##加入CPU和GPU选择
        inputs, target = inputs.to(device), target.to(device)

        optimizer.zero_grad()

        #前馈，反向传播，更新
        outputs = model(inputs)
        loss = model.criterion(outputs, target)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        ##计算每300次打印一次学习效果
        if batch_idx % 300 == 299:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 300))
            running_loss = 0.0


def test():
    correct = 0 
    total = 0
    ##with这里标记是不再计算梯度
    with torch.no_grad():
        for data in test_dataloader:
            inputs, labels = data
            ##加入CPU和GPU选择
            inputs, labels = inputs.to(device), labels.to(device)


            outputs = model(inputs)
            ##预测返回的是两列，第一列是下标就是0-9的值，第二列为预测值，下面的dim=1就是找维度1（第二列）最大值输出
            _, predicted = torch.max(outputs.data, dim=1)

            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    currentpredicted = (100 * correct / total)
    ##用global声明toppredicted,用于在函数内部修改在函数外部声明的全局变量，否则报错
    global toppredicted
    ##当预测率大于原来的保存模型
    if currentpredicted > toppredicted:
        toppredicted = currentpredicted
        torch.save(model.state_dict(), savemodel_name)
        print(savemodel_name+" saved, currentpredicted:%d %%" % currentpredicted)

    print('Accuracy on test set: %d %%' % currentpredicted)        

##开始训练
timestart = time.time()
for epoch in range(10):
    train(epoch)
    test()
timeend = time.time() - timestart
print("use time: {:.0f}m {:.0f}s".format(timeend // 60, timeend % 60))