|
|
@@ -0,0 +1,107 @@
|
|
|
+from model import LeNet5 # 引入我们写好的神经网络（也是我们要训练的网络）
|
|
|
+from torch.nn import CrossEntropyLoss # 引入损失函数（用来做分类损失最佳）
|
|
|
+from torch.optim import Adam # 损失函数优化器（其他损失函数的最小值：使用梯度下降法）
|
|
|
+from torchvision.datasets import MNIST # 数据集
|
|
|
+from torchvision.transforms import Compose, ToTensor # 图像转换为张量
|
|
|
+from torch.utils.data import DataLoader # 批次数据集（数据集不是一次训练，而是分成多个批次）
|
|
|
+import torch # 调用torch的基本函数
|
|
|
+import os # 路径与文件处理
|
|
|
+
|
|
|
+class Trainer:
|
|
|
+ def __init__(self):
|
|
|
+ # 判定电脑安装GPU的环境：cuda,cnn,pytorch-
|
|
|
+ self.CUDA = False # torch.cuda.is_available() # 返回逻辑值：True：支持GPU，False不支持GPU
|
|
|
+ # 准备训练需要的数据、损失函数，优化器，学习率
|
|
|
+ self.lr = 0.0001
|
|
|
+ self.m_file = "lenet5.pt" # 模型的保存文件
|
|
|
+ self.net = LeNet5() # 需要训练的网络
|
|
|
+ if self.CUDA:
|
|
|
+ self.net.cuda() # 把net网络模型存储到GPU上面
|
|
|
+
|
|
|
+ # 加载上一次训练的模型
|
|
|
+ if os.path.exists(self.m_file):
|
|
|
+ # 存在就加载
|
|
|
+ print("加载模型中...")
|
|
|
+ state = torch.load(self.m_file) # 读取文件
|
|
|
+ self.net.load_state_dict(state)
|
|
|
+ else:
|
|
|
+ print("模型存在，重头训练")
|
|
|
+ self.loss_f = CrossEntropyLoss() # 损失函数
|
|
|
+ self.optimizer = Adam(self.net.parameters(), lr=self.lr) # 优化器
|
|
|
+
|
|
|
+
|
|
|
+ # 数据集 - 训练集
|
|
|
+ self.trans = Compose([ToTensor()])
|
|
|
+ self.ds_train = MNIST(root="datasets", download=True, train=True, transform=self.trans)
|
|
|
+ self.bt_train = DataLoader(self.ds_train, batch_size=1000, shuffle=True)
|
|
|
+
|
|
|
+ # 数据集 - 验证集
|
|
|
+ self.ds_valid = MNIST(root="datasets", download=True, train=False, transform=self.trans)
|
|
|
+ self.bt_valid = DataLoader(self.ds_valid, batch_size=1000, shuffle=False) # shuffle=False是否随机打乱
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+ def train_one(self):
|
|
|
+ # 训练一轮epoch，60批次batch，每个批次1000张录像
|
|
|
+ # 循环训练每个批次
|
|
|
+ batch = 1
|
|
|
+ for x, y in self.bt_train:
|
|
|
+ # print(F"\t|-第{batch:02d}批次训练。")
|
|
|
+ if self.CUDA:
|
|
|
+ x = x.cuda()
|
|
|
+ y = y.cuda()
|
|
|
+ y_ = self.net(x) # 进行预测
|
|
|
+ # 计算误差
|
|
|
+ loss = self.loss_f(y_, y) # 使用真实的标签与预测标签计算sunshi
|
|
|
+ # 优化卷积核与全链接矩阵
|
|
|
+ self.optimizer.zero_grad()
|
|
|
+ loss.backward() # 使用导数计算梯度
|
|
|
+ self.optimizer.step() # 更新梯度，优化网络模型
|
|
|
+ batch+=1
|
|
|
+
|
|
|
+ @torch.no_grad() # 该函数中的运算都不会进行求导跟踪
|
|
|
+ def valid(self):
|
|
|
+ # 使用测试数据集验证 (准确率，损失值)
|
|
|
+ all_num = 0.0 # 验证的样本数
|
|
|
+ acc_num = 0.0 # 识别正确数量
|
|
|
+ all_loss = 0.0 # 累加每个样本的损失
|
|
|
+ for t_x, t_y in self.bt_valid:
|
|
|
+ if self.CUDA:
|
|
|
+ t_x = t_x.cuda()
|
|
|
+ t_y = t_y.cuda()
|
|
|
+ # 统计样本数
|
|
|
+ all_num += len(t_x)
|
|
|
+ # 预测
|
|
|
+ t_y_ = self.net(t_x)
|
|
|
+ # 判定准确
|
|
|
+ y_cls = torch.argmax(t_y_, dim=1)
|
|
|
+ # 统计正确率
|
|
|
+ acc_num += (y_cls == t_y).float().sum()
|
|
|
+ # 统计损失
|
|
|
+ all_loss += self.loss_f(t_y_, t_y)
|
|
|
+ print(F"\t|- 训练损失：{all_loss/all_num:8.6f}")
|
|
|
+ print(F"\t|- 准确率：{acc_num * 100.0/all_num:5.2f}%")
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+ def train(self, epoch):
|
|
|
+ # 训练指定的论数
|
|
|
+ for e in range(epoch):
|
|
|
+ print(F"第{e:04d}轮训练。")
|
|
|
+ self.train_one()
|
|
|
+ self.valid()
|
|
|
+ # 保存训练的模型
|
|
|
+ torch.save(self.net.state_dict(), self.m_file)
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ trainer = Trainer()
|
|
|
+ trainer.train(50) # 训练10轮
|
|
|
+ # print(torch.cuda.is_available())
|
|
|
+
|
|
|
+"""
|
|
|
+ 244K = 6 * 5 * 5 * 8 矩阵
|
|
|
+ 16 * 5 * 5 * 8
|
|
|
+ 120 * 5 * 5 * 8
|
|
|
+ 120 * 84* 8
|
|
|
+ 84 * 10*8
|
|
|
+"""
|
