|
@@ -0,0 +1,101 @@
|
|
|
|
|
+
|
|
|
|
|
+import torch
|
|
|
|
|
+import torch.nn as nn
|
|
|
|
|
+from torch.utils.data import DataLoader
|
|
|
|
|
+from torch.utils.tensorboard import SummaryWriter
|
|
|
|
|
+from torchvision import datasets, transforms
|
|
|
|
|
+from net import FullyConnectedNet
|
|
|
|
|
+import torch.nn.functional as F
|
|
|
|
|
+import tqdm #让循环在运行时,自动在控制台显示一个动态更新的进度条
|
|
|
|
|
+
|
|
|
|
|
+# 指定日志目录
|
|
|
|
|
+writer = SummaryWriter(log_dir='./mlp/logs')
|
|
|
|
|
+
|
|
|
|
|
+# 1. 判断是否使用CUDA
|
|
|
|
|
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
|
+print(f'Using device: {device}')
|
|
|
|
|
+# 2. 准备数据,读数据,数据预处理
|
|
|
|
|
+train_set = datasets.MNIST(root='./mlp/data',
|
|
|
|
|
+ train=True, download=True, transform=transforms.ToTensor())
|
|
|
|
|
+test_set = datasets.MNIST(root='./mlp/data',
|
|
|
|
|
+ train=False, download=True, transform=transforms.ToTensor())
|
|
|
|
|
+
|
|
|
|
|
+#显存,CUDA out of memory 错误,可能是因为 batch_size 太大了,导致显存不足。可以尝试减小 batch_size。
|
|
|
|
|
+train_loader = DataLoader(dataset=train_set, batch_size=100, shuffle=True)
|
|
|
|
|
+test_loader = DataLoader(dataset=test_set, batch_size=100, shuffle=False)
|
|
|
|
|
+# 3. 创建模型
|
|
|
|
|
+model = FullyConnectedNet()
|
|
|
|
|
+model = model.to(device)
|
|
|
|
|
+# 4. 确定损失函数
|
|
|
|
|
+loss_fn = nn.MSELoss()
|
|
|
|
|
+# 5. 创建优化器 使用梯度下降算法 参数的更新
|
|
|
|
|
+opt = torch.optim.Adam(model.parameters())
|
|
|
|
|
+
|
|
|
|
|
+# 早停法参数:当测试集损失连续 patience 轮不再变好时,提前终止训练
|
|
|
|
|
+patience = 5
|
|
|
|
|
+best_test_loss = float('inf') # 记录到目前为止最好的测试损失
|
|
|
|
|
+no_improve_count = 0 # 记录测试损失连续没有变好的轮数
|
|
|
|
|
+
|
|
|
|
|
+max_epochs = 1000
|
|
|
|
|
+#轮次, 训练集 60000 张图片,batch_size=100, 60000/100=600
|
|
|
|
|
+for epoch in range(max_epochs):
|
|
|
|
|
+ # 6. 训练模型
|
|
|
|
|
+ model.train()
|
|
|
|
|
+ train_total_loss = 0
|
|
|
|
|
+ for images, labels in tqdm.tqdm(train_loader, desc="train", total=len(train_loader)):
|
|
|
|
|
+ # 将数据移动到设备
|
|
|
|
|
+ images, labels = images.to(device), labels.to(device)
|
|
|
|
|
+ labels = F.one_hot(labels, num_classes=10).float()
|
|
|
|
|
+ outputs = model(images) # 前向传播
|
|
|
|
|
+ loss = loss_fn(outputs, labels) # 计算损失
|
|
|
|
|
+
|
|
|
|
|
+ opt.zero_grad() # 清空梯度
|
|
|
|
|
+ loss.backward() # 反向传播 计算梯度
|
|
|
|
|
+ opt.step() # 更新参数
|
|
|
|
|
+ train_total_loss += loss.item()
|
|
|
|
|
+ train_avg_loss = train_total_loss / len(train_loader)
|
|
|
|
|
+ print(f'Epoch {epoch+1}/{max_epochs}, Loss: {train_avg_loss:.4f}')
|
|
|
|
|
+ # 7. 测试模型
|
|
|
|
|
+ model.eval()
|
|
|
|
|
+ test_total_loss = 0 #一轮测试的总损失
|
|
|
|
|
+ test_total_acc = 0 #一轮测试的总得分
|
|
|
|
|
+ #禁用梯度计算
|
|
|
|
|
+ #在推理时,我们不需要反向传播,因此不需要计算损失函数对参数的梯度
|
|
|
|
|
+ with torch.inference_mode():
|
|
|
|
|
+ for images, labels in tqdm.tqdm(test_loader, desc="test", total=len(test_loader)):
|
|
|
|
|
+ images, labels = images.to(device), labels.to(device)
|
|
|
|
|
+ labels = F.one_hot(labels, num_classes=10).float()
|
|
|
|
|
+ outputs = model(images)
|
|
|
|
|
+ # 计算损失
|
|
|
|
|
+ loss = loss_fn(outputs, labels)
|
|
|
|
|
+ test_total_loss += loss.item()
|
|
|
|
|
+
|
|
|
|
|
+ pred = torch.argmax(outputs, dim=1)
|
|
|
|
|
+ target = torch.argmax(labels, dim=1)
|
|
|
|
|
+ acc = torch.eq(pred, target).float().mean()
|
|
|
|
|
+ test_total_acc += acc.item()
|
|
|
|
|
+
|
|
|
|
|
+ test_avg_acc = test_total_acc / len(test_loader)
|
|
|
|
|
+ print(f'epoch:{epoch+1}, Test Accuracy: {test_avg_acc:.4f}')
|
|
|
|
|
+
|
|
|
|
|
+ test_avg_loss = test_total_loss / len(test_loader)
|
|
|
|
|
+ print(f'epoch:{epoch+1},Test Loss: {test_avg_loss:.4f}')
|
|
|
|
|
+ # 记录数据
|
|
|
|
|
+ writer.add_scalars('Loss/train', {'train_avg_loss': train_avg_loss,
|
|
|
|
|
+ 'test_avg_loss': test_avg_loss}, epoch)
|
|
|
|
|
+ writer.add_scalar('Accuracy/test', test_avg_acc, epoch)
|
|
|
|
|
+ # 8. 早停判断与模型保存
|
|
|
|
|
+ if test_avg_loss < best_test_loss:
|
|
|
|
|
+ # 测试损失变好了:更新最好成绩,重置计数器,保存当前最优模型
|
|
|
|
|
+ best_test_loss = test_avg_loss
|
|
|
|
|
+ no_improve_count = 0
|
|
|
|
|
+ torch.save(model.state_dict(), './mlp/model/mnist_net_best.pth')
|
|
|
|
|
+ print(f'测试损失改善,保存最优模型 (best loss: {best_test_loss:.4f})')
|
|
|
|
|
+ else:
|
|
|
|
|
+ # 测试损失没有变好:计数器加一
|
|
|
|
|
+ no_improve_count += 1
|
|
|
|
|
+ print(f'测试损失未改善 ({no_improve_count}/{patience})')
|
|
|
|
|
+ if no_improve_count >= patience:
|
|
|
|
|
+ # 连续 patience 轮没有变好,提前终止训练
|
|
|
|
|
+ print(f'早停:测试损失连续 {patience} 轮未改善,在第 {epoch+1} 轮终止训练')
|
|
|
|
|
+ break
|