本帖最后由 knv 于 2024-5-8 18:10 编辑
补充准确率: 验证准确率 99.16%
新版本代码:
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
class OptimizedConvNet(nn.Module):
def __init__(self,num_classes=10):
super(OptimizedConvNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.layer2 = nn.Sequential(
nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.drop_out = nn.Dropout()
self.fc1 = nn.Linear(7 * 7 * 64, 1000)
self.fc2 = nn.Linear(1000, num_classes)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = out.reshape(out.size(0), -1)
out = self.drop_out(out)
out = self.fc1(out)
out = self.fc2(out)
return out
# 定义超参数
learning_rate = 0.001
batch_size = 64
num_epochs = 10
# 检查是否有可用的GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 创建模型,并将模型移动到GPU上
model = OptimizedConvNet(num_classes=10).to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# 加载MNIST数据集
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# 加载验证集
valid_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
valid_loader = DataLoader(valid_dataset, batch_size=64, shuffle=False)
# 初始化最优验证损失为无穷大
best_valid_loss = float('inf')
import matplotlib.pyplot as plt
# 初始化准确率列表
train_accuracies = []
valid_accuracies = []
for epoch in range(num_epochs):
# 在训练集上训练
model.train()
total_train_loss = 0
correct_train_preds = 0
total_train_preds = 0
for i, (images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
_, predicted = torch.max(outputs.data, 1)
total_train_preds += labels.size(0)
correct_train_preds += (predicted == labels).sum().item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0:
print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(train_loader)}], Loss: {loss.item()}')
total_train_loss += loss.item()
average_train_loss = total_train_loss / len(train_loader)
train_accuracy = correct_train_preds / total_train_preds
train_accuracies.append(train_accuracy)
print(f'Epoch: {epoch + 1}, Training Loss: {average_train_loss:.4f}, Training Accuracy: {train_accuracy:.4f}')
# 在验证集上验证
model.eval()
valid_loss = 0.0
correct_valid_preds = 0
total_valid_preds = 0
with torch.no_grad():
for images, labels in valid_loader:
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
_, predicted = torch.max(outputs.data, 1)
total_valid_preds += labels.size(0)
correct_valid_preds += (predicted == labels).sum().item()
valid_loss += loss.item()
valid_loss /= len(valid_loader)
valid_accuracy = correct_valid_preds / total_valid_preds
valid_accuracies.append(valid_accuracy)
print(f'Epoch: {epoch+1}, Validation Loss: {valid_loss:.4f}, Validation Accuracy: {valid_accuracy:.4f}')
# 如果验证损失有所下降,则保存模型
if valid_loss < best_valid_loss:
best_valid_loss = valid_loss
torch.save(model.state_dict(), 'best_model.pth')
print('Model saved.')
# 绘制训练和验证准确率
plt.plot(range(1, num_epochs + 1), train_accuracies, label='Train')
plt.plot(range(1, num_epochs + 1), valid_accuracies, label='Valid')
plt.title('Model Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.legend()
plt.show()