<code>import torch

import torch.nn as nn

import torch.optim as optim

import torchvision

import torchvision.transforms as transforms

import matplotlib.pyplot as plt

# ...（之前的代码保持不变，包括设备配置、数据加载、模型定义等）

# 训练模型

def train(model, device, train_loader, optimizer, epoch, num_epochs):

model.train() # Set the model to training mode

total_step = len(train_loader)

losses = []

for i, (images, labels) in enumerate(train_loader):

images, labels = images.to(device), labels.to(device)

# Forward pass

outputs = model(images)

loss = criterion(outputs, labels)

# Backward and optimize

optimizer.zero_grad()

loss.backward()

optimizer.step()

# Collect loss for plotting

losses.append(loss.item())

if (i+1) % 100 == 0:

print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'

.format(epoch+1, num_epochs, i+1, total_step, loss.item()))

return losses

# 绘制损失曲线的函数

def plot_losses(epochs, losses):

plt.figure(figsize=(10, 5))

for i, loss_per_epoch in enumerate(losses):

plt.plot(loss_per_epoch, label=f'Epoch { i+1}')

plt.title('Loss over epochs')

plt.xlabel('Steps')

plt.ylabel('Loss')

plt.legend()

plt.show()

# 训练过程

losses_over_epochs = []

for epoch in range(num_epochs):

losses = train(model, device, train_loader, optimizer, epoch, num_epochs)

losses_over_epochs.append(losses)

print(f'Epoch { epoch+1}/{ num_epochs} - Average Loss: { sum(losses)/len(losses):.4f}')

# 绘制所有epoch的损失曲线

plot_losses(num_epochs, losses_over_epochs)

# ...（测试模型和保存模型的代码保持不变）

在这个修改后的代码中，我们添加了两个新的函数：

train：这个函数用于训练模型，并记录每个step的损失。它返回一个包含所有step损失的列表。plot_losses：这个函数接受epoch列表和损失列表作为参数，并绘制出损失随训练step变化的曲线。

在主训练循环中，我们对每个epoch调用train函数，并收集所有epoch的损失，然后使用plot_losses函数绘制损失曲线。

请注意，这里绘制的是每个step的损失，而不是每个epoch的损失均值。如果想要绘制每个epoch的平均损失，可以修改train函数来计算每个epoch的平均损失，并只记录这个值。

咱直接把最终代码贴上：

import torch

import torch.nn as nn

import torch.optim as optim

import torchvision

import torchvision.transforms as transforms

import matplotlib.pyplot as plt

# Device configuration

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Hyper parameters

num_epochs = 10

num_classes = 10

batch_size = 100

learning_rate = 0.001

# MNIST dataset

train_dataset = torchvision.datasets.MNIST(root='../../data/',code>

train=True,

transform=transforms.ToTensor(),

download=True)

test_dataset = torchvision.datasets.MNIST(root='../../data/',code>

train=False,

transform=transforms.ToTensor())

# Data loader

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,

batch_size=batch_size,

shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,

batch_size=batch_size,

shuffle=False)

# Convolutional neural network (two convolutional layers)

class ConvNet(nn.Module):

def __init__(self, num_classes=10):

super(ConvNet, self).__init__()

self.layer1 = nn.Sequential(

nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),

nn.BatchNorm2d(16),

nn.ReLU(),

nn.MaxPool2d(kernel_size=2, stride=2))

self.layer2 = nn.Sequential(

nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),

nn.BatchNorm2d(32),

nn.ReLU(),

nn.MaxPool2d(kernel_size=2, stride=2))

self.fc = nn.Linear(7*7*32, num_classes)

def forward(self, x):

out = self.layer1(x)

out = self.layer2(out)

out = out.view(out.size(0), -1) # Flatten the output

out = self.fc(out)

return out

model = ConvNet(num_classes).to(device)

# Loss and optimizer

criterion = nn.CrossEntropyLoss()

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

# Function to visualize images

def show_images(images):

plt.figure(figsize=(10,10))

for i, img in enumerate(images):

plt.subplot(5, 5, i+1)

plt.xticks([])

plt.yticks([])

plt.grid(False)

plt.imshow(img.squeeze().numpy(), cmap='gray')code>

plt.show()

# Function to train the model

def train(model, device, train_loader, optimizer, epoch, num_epochs):

model.train() # Set the model to training mode

total_step = len(train_loader)

losses = []

for i, (images, labels) in enumerate(train_loader):

images, labels = images.to(device), labels.to(device)

optimizer.zero_grad()

outputs = model(images)

loss = criterion(outputs, labels)

loss.backward()

optimizer.step()

losses.append(loss.item())

if (i+1) % 100 == 0:

print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'

.format(epoch+1, num_epochs, i+1, total_step, loss.item()))

return losses

# Function to test the model

def test(model, device, test_loader):

model.eval() # Set the model to evaluation mode

correct = 0

total = 0

with torch.no_grad():

for images, labels in test_loader:

images = images.to(device)

labels = labels.to(device)

outputs = model(images)

_, predicted = torch.max(outputs.data, 1)

total += labels.size(0)

correct += (predicted == labels).sum().item()

return 100 * correct / total

# Function to plot training progress

def plot_progress(epochs, train_losses, test_accuracies):

plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)

for i in range(epochs):

plt.plot(train_losses[i], label=f'Epoch { i+1}')

plt.title('Training Loss')

plt.xlabel('Batch')

plt.ylabel('Loss')

plt.legend()

plt.subplot(1, 2, 2)

plt.plot(test_accuracies, label='Accuracy')code>

plt.title('Test Accuracy')

plt.xlabel('Epoch')

plt.ylabel('Accuracy (%)')

plt.legend()

plt.show()

# Visualize a few images

dataiter = iter(train_loader)

images, _ = next(dataiter)

show_images(images[:25]) # Visualize 25 images

# Initialize lists to monitor loss and accuracy

train_losses = []

test_accuracies = []

# Train the model

for epoch in range(num_epochs):

print(f'Epoch { epoch+1}/{ num_epochs}')

train_loss = train(model, device, train_loader, optimizer, epoch, num_epochs)

train_losses.append(train_loss)

test_accuracy = test(model, device, test_loader)

test_accuracies.append(test_accuracy)

print(f'Epoch { epoch+1}/{ num_epochs} - Average Loss: { sum(train_loss)/len(train_loss):.4f}, Accuracy: { test_accuracy:.2f}%')

# Plot training progress

plot_progress(num_epochs, train_losses, test_accuracies)

# Save the model checkpoint

torch.save(model.state_dict(), 'model.ckpt')

经过修改后的训练过程如下图所示：

还有很多很多小改进的方向，就留给各位自己尝试啦。