|
|
@@ -0,0 +1,74 @@
|
|
|
+
|
|
|
+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 lenet import LeNet5
|
|
|
+import tqdm
|
|
|
+
|
|
|
+
|
|
|
+# 指定日志目录
|
|
|
+writer = SummaryWriter(log_dir='D:/gitcode/cnn_dropout_bn/logs')
|
|
|
+
|
|
|
+# 1. 判断是否使用CUDA
|
|
|
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
+print(f'Using device: {device}')
|
|
|
+# 2. 准备数据
|
|
|
+train_set = datasets.CIFAR10(root='D:/gitcode/cnn_dropout_bn/data', train=True, download=True, transform=transforms.ToTensor())
|
|
|
+test_set = datasets.CIFAR10(root='D:/gitcode/cnn_dropout_bn/data', train=False, download=True, transform=transforms.ToTensor())
|
|
|
+train_loader = DataLoader(dataset=train_set, batch_size=128, shuffle=True)
|
|
|
+test_loader = DataLoader(dataset=test_set, batch_size=128, shuffle=False)
|
|
|
+# 3. 创建模型
|
|
|
+model = LeNet5()
|
|
|
+model = model.to(device)
|
|
|
+# 4. 确定损失函数
|
|
|
+loss_fn = nn.CrossEntropyLoss()
|
|
|
+# 5. 创建优化器 使用梯度下降算法 参数的更新
|
|
|
+opt = torch.optim.Adam(model.parameters(),lr=0.0003)
|
|
|
+
|
|
|
+for epoch in range(100):
|
|
|
+ # 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)
|
|
|
+ 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}, Loss: {train_avg_loss:.4f}')
|
|
|
+ # 7. 测试模型
|
|
|
+ model.eval()# 关闭dropout, batchsize=1, 5
|
|
|
+ #10000
|
|
|
+ #模型在做推理时, BN使用的是整个训练集上均值和方差,而不是当前batch的均值和方差
|
|
|
+ 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)
|
|
|
+ outputs = model(images)
|
|
|
+ # 计算损失
|
|
|
+ loss = loss_fn(outputs, labels)
|
|
|
+ test_total_loss += loss.item()
|
|
|
+
|
|
|
+ pred = torch.argmax(outputs, dim=1)
|
|
|
+ acc = torch.eq(pred, labels).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. 保存模型
|
|
|
+ torch.save(model.state_dict(), 'D:/gitcode/cnn_dropout_bn/model/lenet5.pth')
|