111 lines
3.5 KiB
Python
111 lines
3.5 KiB
Python
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import torch.optim as optim
|
|||
|
|
from torchvision import datasets, transforms
|
|||
|
|
import matplotlib.pyplot as plt
|
|||
|
|
|
|||
|
|
# 1. 数据准备（以MNIST手写数字识别为例）
|
|||
|
|
transform = transforms.Compose([
|
|||
|
|
transforms.ToTensor(),
|
|||
|
|
transforms.Normalize((0.5,), (0.5,)) # 像素值归一化到[-1,1]
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
train_set = datasets.MNIST('data', download=True, train=True, transform=transform)
|
|||
|
|
test_set = datasets.MNIST('data', download=True, train=False, transform=transform)
|
|||
|
|
|
|||
|
|
train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)
|
|||
|
|
test_loader = torch.utils.data.DataLoader(test_set, batch_size=1000)
|
|||
|
|
|
|||
|
|
# 2. 神经网络模型（演示梯度控制技巧）
|
|||
|
|
class Net(nn.Module):
|
|||
|
|
def __init__(self):
|
|||
|
|
super(Net, self).__init__()
|
|||
|
|
self.fc1 = nn.Linear(784, 128)
|
|||
|
|
self.fc2 = nn.Linear(128, 64)
|
|||
|
|
self.fc3 = nn.Linear(64, 10)
|
|||
|
|
|
|||
|
|
# He初始化适配ReLU
|
|||
|
|
nn.init.kaiming_normal_(self.fc1.weight, nonlinearity='relu')
|
|||
|
|
nn.init.kaiming_normal_(self.fc2.weight, nonlinearity='relu')
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
x = x.view(-1, 784) # 展平图像
|
|||
|
|
x = torch.relu(self.fc1(x))
|
|||
|
|
x = torch.relu(self.fc2(x))
|
|||
|
|
x = self.fc3(x) # 输出层无需激活（CrossEntropyLoss内置Softmax）
|
|||
|
|
return x
|
|||
|
|
|
|||
|
|
# 3. 训练配置
|
|||
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|||
|
|
model = Net().to(device)
|
|||
|
|
optimizer = optim.Adam(model.parameters(), lr=0.001)
|
|||
|
|
criterion = nn.CrossEntropyLoss()
|
|||
|
|
|
|||
|
|
# 梯度裁剪阈值
|
|||
|
|
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=2.0)
|
|||
|
|
|
|||
|
|
# 4. 训练过程可视化记录
|
|||
|
|
train_losses = []
|
|||
|
|
accuracies = []
|
|||
|
|
|
|||
|
|
def train(epoch):
|
|||
|
|
model.train()
|
|||
|
|
for batch_idx, (data, target) in enumerate(train_loader):
|
|||
|
|
data, target = data.to(device), target.to(device)
|
|||
|
|
optimizer.zero_grad()
|
|||
|
|
output = model(data)
|
|||
|
|
loss = criterion(output, target)
|
|||
|
|
loss.backward()
|
|||
|
|
optimizer.step()
|
|||
|
|
|
|||
|
|
# 记录训练损失
|
|||
|
|
if batch_idx % 100 == 0:
|
|||
|
|
train_losses.append(loss.item())
|
|||
|
|
|
|||
|
|
# 5. 测试函数（含准确率计算）
|
|||
|
|
def test():
|
|||
|
|
model.eval()
|
|||
|
|
test_loss = 0
|
|||
|
|
correct = 0
|
|||
|
|
with torch.no_grad():
|
|||
|
|
for data, target in test_loader:
|
|||
|
|
data, target = data.to(device), target.to(device)
|
|||
|
|
output = model(data)
|
|||
|
|
test_loss += criterion(output, target).item()
|
|||
|
|
pred = output.argmax(dim=1, keepdim=True)
|
|||
|
|
correct += pred.eq(target.view_as(pred)).sum().item()
|
|||
|
|
|
|||
|
|
accuracy = 100. * correct / len(test_loader.dataset)
|
|||
|
|
accuracies.append(accuracy)
|
|||
|
|
return test_loss
|
|||
|
|
|
|||
|
|
# 6. 执行训练（3个epoch演示）
|
|||
|
|
for epoch in range(1, 4):
|
|||
|
|
train(epoch)
|
|||
|
|
loss = test()
|
|||
|
|
print(f'Epoch {epoch}: Test Loss={loss:.4f}, Accuracy={accuracies[-1]:.2f}%')
|
|||
|
|
|
|||
|
|
# 7. 可视化训练过程
|
|||
|
|
plt.figure(figsize=(12,5))
|
|||
|
|
plt.subplot(1,2,1)
|
|||
|
|
plt.plot(train_losses, label='Training Loss')
|
|||
|
|
plt.title("Loss Curve")
|
|||
|
|
plt.subplot(1,2,2)
|
|||
|
|
plt.plot(accuracies, label='Accuracy', color='orange')
|
|||
|
|
plt.title("Accuracy Curve")
|
|||
|
|
plt.show()
|
|||
|
|
|
|||
|
|
# 8. 示例预测展示
|
|||
|
|
sample_data, sample_label = next(iter(test_loader))
|
|||
|
|
with torch.no_grad():
|
|||
|
|
prediction = model(sample_data.to(device)).argmax(dim=1)
|
|||
|
|
|
|||
|
|
# 显示预测结果对比
|
|||
|
|
plt.figure(figsize=(10,6))
|
|||
|
|
for i in range(6):
|
|||
|
|
plt.subplot(2,3,i+1)
|
|||
|
|
plt.imshow(sample_data[i][0], cmap='gray')
|
|||
|
|
plt.title(f"True: {sample_label[i]}\nPred: {prediction[i].item()}")
|
|||
|
|
plt.axis('off')
|
|||
|
|
plt.tight_layout()
|
|||
|
|
plt.show()
|