import torch  # 深度學(xué)習(xí)框架核心
import torch.nn as nn  # 神經(jīng)網(wǎng)絡(luò)模塊
from torchvision import datasets, transforms  # 數(shù)據(jù)處理利器

transform = transforms.Compose([
    transforms.ToTensor(),  # 將圖像轉(zhuǎn)換為數(shù)值矩陣
    transforms.Normalize((0.1307,), (0.3081,))  # 標(biāo)準(zhǔn)化處理
])

plt.imshow(images[0].squeeze(), cmap='gray')
plt.title(f'Label: {labels[0]}')

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()  # 將圖片展開為向量
        self.fc1 = nn.Linear(28*28, 128)  # 第一隱藏層
        self.fc2 = nn.Linear(128, 64)    # 第二隱藏層
        self.fc3 = nn.Linear(64, 10)     # 輸出層
        self.dropout = nn.Dropout(0.5)   # 正則化裝置

def forward(self, x):
    x = self.flatten(x)  # 展平操作：將圖片變?yōu)?84維向量
    x = torch.relu(self.fc1(x))  # 通過(guò)第一個(gè)全連接層
    x = self.dropout(x)         # 隨機(jī)屏蔽部分神經(jīng)元
    x = torch.relu(self.fc2(x))  # 第二個(gè)全連接層
    return self.fc3(x)          # 最終輸出10個(gè)數(shù)字的概率

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

def train(epoch):
    model.train()  # 切換至訓(xùn)練模式
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()  # 清空之前的梯度
        output = model(data)    # 前向傳播
        loss = criterion(output, target)  # 計(jì)算損失值
        loss.backward()         # 反向傳播求梯度
        optimizer.step()        # 更新網(wǎng)絡(luò)參數(shù)

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision

from torchvision import transforms
import matplotlib.pyplot as plt

# 2. 數(shù)據(jù)準(zhǔn)備
# 定義數(shù)據(jù)預(yù)處理：轉(zhuǎn)換為Tensor并標(biāo)準(zhǔn)化（MNIST的均值和標(biāo)準(zhǔn)差）
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加載訓(xùn)練集和測(cè)試集
train_dataset = torchvision.datasets.MNIST(
    root='./data', 
    train=True,
    download=True,
    transform=transform
)

test_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    transform=transform
)

# 創(chuàng)建數(shù)據(jù)加載器
train_loader = torch.utils.data.DataLoader(
    train_dataset,

    batch_size=64,
    shuffle=True
)

test_loader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=1000,
    shuffle=False
)

# 查看數(shù)據(jù)集信息
print(f'Train samples: {len(train_dataset)}')
print(f'Test samples: {len(test_dataset)}')

# 可視化樣本
images, labels = next(iter(train_loader))
plt.imshow(images[0].squeeze(), cmap='gray')
plt.title(f'Label: {labels[0]}')
plt.show()

# 3. 定義神經(jīng)網(wǎng)絡(luò)模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)
        self.dropout = nn.Dropout(0.5)
        
    def forward(self, x):
        x = self.flatten(x)
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = Net()
print(model)

# 4. 定義損失函數(shù)和優(yōu)化器
criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=0.001)


# 5. 訓(xùn)練模型
def train(epoch):

    model.train()
    running_loss = 0.0
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        

        running_loss += loss.item()
        if batch_idx % 100 == 99:

            print(f'Epoch: {epoch+1}, Batch: {batch_idx+1}, Loss: {running_loss/100:.3f}')
            running_loss = 0.0

# 6. 測(cè)試模型

def test():
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            outputs = model(data)
            _, predicted = torch.max(outputs.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    
    accuracy = 100 * correct / total
    print(f'Test Accuracy: {accuracy:.2f}%')
    return accuracy

# 7. 執(zhí)行訓(xùn)練和測(cè)試

epochs = 5
for epoch in range(epochs):
    train(epoch)
    test()

# 8. 保存模型
torch.save(model.state_dict(), 'mnist_model.pth')

# 擴(kuò)展測(cè)試函數(shù)，增加按數(shù)字統(tǒng)計(jì)的功能

def detailed_test():
    model.eval()
    class_correct = [0] * 10  # 存儲(chǔ)每個(gè)數(shù)字的正確計(jì)數(shù)
    class_total = [0] * 10    # 存儲(chǔ)每個(gè)數(shù)字的總樣本數(shù)
    
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            
            # 遍歷每個(gè)預(yù)測(cè)結(jié)果
            for label, prediction in zip(labels, predicted):
                class_total[label] += 1
                if label == prediction:
                    class_correct[label] += 1


    # 打印每個(gè)數(shù)字的準(zhǔn)確率
    print("{:^10} | {:^10} | {:^10}".format("數(shù)字", "正確數(shù)", "準(zhǔn)確率"))
    print("-"*33)

    for i in range(10):
        acc = 100 * class_correct[i] / class_total[i]
        print("{:^10} | {:^10} | {:^10.2f}%".format(i, class_correct[i], acc))
    
    # 可視化錯(cuò)誤案例
    wrong_examples = []

    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        mask = predicted != labels
        wrong_examples.extend(zip(images[mask], labels[mask], predicted[mask]))
    
    # 隨機(jī)展示3個(gè)錯(cuò)誤樣本

    fig, axes = plt.subplots(1, 3, figsize=(12,4))

    for ax, (img, true, pred) in zip(axes, wrong_examples[:3]):
        ax.imshow(img.squeeze(), cmap='gray')

        ax.set_title(f'True: {true}\nPred: {pred}')
        ax.axis('off')
    plt.show()

# 執(zhí)行詳細(xì)測(cè)試
detailed_test()

# PyTorch版
import torch.nn as nn
layer = nn.Linear(in_features=784, out_features=128)

# TensorFlow版
from tensorflow.keras.layers import Dense
layer = Dense(units=128, input_dim=784)

for epoch in range(epochs):
    for data, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
model.fit(train_dataset, epochs=epochs)  # 自動(dòng)完成訓(xùn)練循環(huán)