pip install torch torchvision matplotlib

import torch
import torch.nn as nn  # 用于構(gòu)建神經(jīng)網(wǎng)絡(luò)
import torch.optim as optim  # 用于優(yōu)化網(wǎng)絡(luò)
import torchvision  # 包含了流行的數(shù)據(jù)集和模型
import torchvision.transforms as transforms  # 用于數(shù)據(jù)增強(qiáng)和預(yù)處理
import matplotlib.pyplot as plt  # 用于繪圖和數(shù)據(jù)可視化

transform = transforms.Compose([
    transforms.Resize((32, 32)),  # 將所有圖像調(diào)整為32x32像素
    transforms.ToTensor(),  # 將圖像轉(zhuǎn)換為Tensor格式，范圍為[0, 1]
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 標(biāo)準(zhǔn)化到[-1, 1]范圍
])

# 加載訓(xùn)練集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

# 加載測(cè)試集
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)

# CIFAR-10數(shù)據(jù)集中的類別
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)  # 第一層卷積，輸入通道3（RGB），輸出通道6，卷積核大小5x5
        self.pool = nn.MaxPool2d(2, 2)  # 最大池化層，窗口大小2x2
        self.conv2 = nn.Conv2d(6, 16, 5)  # 第二層卷積，輸入通道6，輸出通道16，卷積核大小5x5
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 全連接層，輸入維度16*5*5，輸出維度120
        self.fc2 = nn.Linear(120, 84)  # 第二個(gè)全連接層，輸入維度120，輸出維度84
        self.fc3 = nn.Linear(84, 10)  # 最后一層，全連接層，輸出維度10（對(duì)應(yīng)CIFAR-10的10個(gè)類別）

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  # 卷積 -> ReLU激活 -> 最大池化
        x = self.pool(F.relu(self.conv2(x)))  # 卷積 -> ReLU激活 -> 最大池化
        x = x.view(-1, 16 * 5 * 5)  # 展平操作，將卷積層的輸出展平成一維向量
        x = F.relu(self.fc1(x))  # 全連接 -> ReLU激活
        x = F.relu(self.fc2(x))  # 全連接 -> ReLU激活
        x = self.fc3(x)  # 全連接層輸出分類結(jié)果
        return x

net = Net()

criterion = nn.CrossEntropyLoss()  # 交叉熵?fù)p失，用于分類任務(wù)
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)  # 隨機(jī)梯度下降優(yōu)化器，帶動(dòng)量

for epoch in range(2):  # 訓(xùn)練2個(gè)epoch
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data  # 獲取輸入數(shù)據(jù)和對(duì)應(yīng)的標(biāo)簽

        optimizer.zero_grad()  # 清零梯度緩存

        outputs = net(inputs)  # 前向傳播：計(jì)算輸出
        loss = criterion(outputs, labels)  # 計(jì)算損失
        loss.backward()  # 反向傳播：計(jì)算梯度
        optimizer.step()  # 更新模型參數(shù)

        running_loss += loss.item()
        if i % 2000 == 1999:  # 每2000個(gè)mini-batch打印一次損失
            print(f'[Epoch {epoch + 1}, Mini-batch {i + 1}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0

print('Finished Training')

correct = 0
total = 0
with torch.no_grad():  # 禁用梯度計(jì)算，以節(jié)省內(nèi)存和加速計(jì)算
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)  # 獲取最大值的索引，即預(yù)測(cè)的類別
        total += labels.size(0)  # 累計(jì)樣本總數(shù)
        correct += (predicted == labels).sum().item()  # 累計(jì)正確預(yù)測(cè)的樣本數(shù)

print(f'Accuracy of the network on the 10000 test images: {100 * correct / total}%')