使用 pytorch 構建一個簡單的卷積神經網絡(CNN)模型,完成對 CIFAR-10 數據集的圖像分類任務。 直接使用 CNN 進行分類的模型性能。 提示: 數據集:CIFAR-10 網絡結構:可以使用 2-3 層卷積層,ReLU 激活,MaxPooling 層,最后連接全連接層。
#1. 數據預處理與加載
import torch
import torchvision
import torchvision.transforms as transforms# 數據增強與歸一化(使用CIFAR-10官方均值和標準差)
transform_train = transforms.Compose([transforms.RandomCrop(32, padding=4), # 隨機裁剪增強泛化性transforms.RandomHorizontalFlip(), # 隨機水平翻轉transforms.ToTensor(),transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])transform_test = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])# 加載數據集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)# 數據加載器
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)#2. CNN模型架構
import torch.nn as nn
import torch.nn.functional as Fclass SimpleCNN(nn.Module):def __init__(self):super(SimpleCNN, self).__init__()self.conv1 = nn.Conv2d(3, 32, 3, padding=1) # 輸入通道3(RGB),輸出32通道self.bn1 = nn.BatchNorm2d(32) # 批量歸一化self.conv2 = nn.Conv2d(32, 64, 3, padding=1)self.bn2 = nn.BatchNorm2d(64)self.conv3 = nn.Conv2d(64, 128, 3, padding=1)self.bn3 = nn.BatchNorm2d(128)self.pool = nn.MaxPool2d(2, 2) # 池化層(尺寸減半)self.fc1 = nn.Linear(128 * 4 * 4, 256) # 全連接層(輸入尺寸計算:32x32 → 16x16 → 8x8 → 4x4)self.fc2 = nn.Linear(256, 10) # 輸出10類def forward(self, x):x = self.pool(F.relu(self.bn1(self.conv1(x)))) # 32x32 → 16x16x = self.pool(F.relu(self.bn2(self.conv2(x)))) # 16x16 → 8x8x = self.pool(F.relu(self.bn3(self.conv3(x)))) # 8x8 → 4x4x = x.view(-1, 128 * 4 * 4) # 展平x = F.relu(self.fc1(x))x = self.fc2(x)return x# 實例化模型并移至GPU(若可用)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = SimpleCNN().to(device)#3. 訓練與優化
import torch.optim as optimcriterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1) # 每5輪學習率×0.1# 訓練循環(10個epoch)
for epoch in range(10):net.train()running_loss = 0.0for i, (inputs, labels) in enumerate(trainloader):inputs, labels = inputs.to(device), labels.to(device)optimizer.zero_grad()outputs = net(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()if i % 100 == 99: # 每100批次打印一次print(f'Epoch [{epoch+1}/10], Step [{i+1}/{len(trainloader)}], Loss: {running_loss/100:.3f}')running_loss = 0.0scheduler.step() # 更新學習率print(f"Epoch {epoch+1} completed, learning rate: {scheduler.get_last_lr()[0]:.6f}")#4. 模型評估與可視化
net.eval()
correct, total = 0, 0
with torch.no_grad():for (images, labels) in testloader:images, labels = images.to(device), labels.to(device)outputs = net(images)_, predicted = torch.max(outputs.data, 1)total += labels.size(0)correct += (predicted == labels).sum().item()accuracy = 100 * correct / total
print(f'Test Accuracy: {accuracy:.2f}%')
運行結果: