import torch  # 引入模組PyTorch
from torchvision import transforms  # 從torch視覺中引入轉換函式
from torchvision import datasets  # 匯入資料庫
from torch.utils.data import DataLoader  # 匯入資料載入器
import torch.nn.functional as F  # 匯入啟用函式
import torch.optim as optim  # 匯入優化器

# 使用的數字分類資料集，只有一個灰度通道0-255
batch_size = 64
transform = transforms.Compose([
    transforms.ToTensor(),  # 將影象的通道從[0,255]對映到[0,1]
    transforms.Normalize((0.1307,), (0.3081,))  # 根據手寫數字集的平均值和標準差來實現歸一化
])  # 定義mnist資料集轉化管道
train_dataset = datasets.MNIST(root='./dataset/mnist/',  # 設定資料集存放的根目錄
                               train=True,  # 選擇是否是訓練集
                               download=True,  # 是否從網上下載
                               transform=transform)  # 使用可選引數trannsform來使用前面定義的轉換
train_loader = DataLoader(dataset=train_dataset,  # 使用資料載入器來讀入訓練集資料
                          shuffle=True,  # 使用可選引數shuffle=True來打亂訓練集
                          batch_size=batch_size)  # 使用可選引數batch_size來確定每批的資料數量
test_dataset = datasets.MNIST(root='./dataset/mnist/',
                              train=False,  # train=False 在這裡代表選用了測試集
                              download=True,
                              transform=transform)
test_loader = DataLoader(dataset=test_dataset,
                         shuffle=False,  # 測試機資料不用打亂
                         batch_size=batch_size)  # 每批資料數量同前面一樣


class Net(torch.nn.Module):  # 建立網路類
    def __init__(self):  # 初始化類
        super(Net, self).__init__()  # 繼承父類
        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)

    def forward(self, x):
        x = x.view(-1, 784)  # 將輸入的shape=(28,28)的張量展平為(batch_size,28*28)
        x = F.gelu(self.l1(x))  # 將輸入通過線性層後再經由啟用函式gelu進入到下一線性層
        x = F.gelu(self.l2(x))
        x = F.gelu(self.l3(x))
        x = F.gelu(self.l4(x))
        return self.l5(x)


model = Net()  # 將神經網路模型例項化給model
criterion = torch.nn.CrossEntropyLoss()  # 定義多分類損失函式
# 使用SGD(Stochastic Gradient for Tensor Decomposition)隨機梯度下降優化器來修正模型的引數
# 並定義了學習率和動量來加速網路的收斂
optimizer = torch.optim.SGD(model.parameters(), lr=0.01,
                            momentum=.5)


def train(epoch):  # 定義訓練函式
    running_loss = 0  # 初始化損失為0
    for batch_idx, (x, y) in enumerate(train_loader):  # 從可迭代物件train_loader中獲取獲取批數，還有輸入和輸出
        inputs, target = x, y  # 將x，y分別給輸入和目標值
        optimizer.zero_grad()  # 初始化將梯度置0
        outputs = model(inputs)  # 輸入經由model的正向傳播得到輸入
        loss = criterion(outputs, target)  # 通過criterion函式來計算損失
        loss.backward()  # 將損失函式進行反向傳播計算梯度
        optimizer.step()  # 根據梯度來更新引數
        running_loss += loss  # 計算每300批的損失值
        if batch_idx % 300 == 299:
            print('[%d,%5d] loss:%.3f' % (epoch + 1, batch_idx + 1, running_loss / 300))  # 輸出300批損失值的平均值
            running_loss = 0  # 將損失值置0


def test():  # 定義測試集
    corrent = 0  # 定義正確率
    total = 0  # 定義總數
    with torch.no_grad():  # 因為是測試集，不用計算跟蹤梯度
        for data in test_loader:
            images, labels = data
            outputs = model(images)
            _, prediction = torch.max(outputs, dim=1)  # torch.max()函式按維度dim返回最大值的那個元素和索引
            total += labels.size(0)  # labels的size(0)就是每批資料的數目
            corrent += (prediction == labels).sum().item()  # 用預測正確的數目的和除以總數目來獲得正確率
    print('Accuracy on test set %d %%' % (100 * corrent / total))


if __name__ == '__main__':
    for epoch in range(10):  # 訓練十輪
        train(epoch)
        test()
 

執行結果：