use ignite for exp2

exp2/modified/train.py

@@ -6,188 +6,184 @@ import torchinfo
 import ignite.engine
 import ignite.metrics
 from ignite.engine import Engine, Events
+from ignite.metrics import Accuracy, Loss
 from ignite.handlers.tqdm_logger import ProgressBar
-from dataset import MnistDataSource
+from dataset import MnistDataLoaders
 from model import Cnn
+import settings
 
 sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
-import pytorch_gpu_utils
+import gpu_utils
 
 
-class Trainer:
-    N_EPOCH: typing.ClassVar[int] = 5
-    N_BATCH_SIZE: typing.ClassVar[int] = 1000
-
-    device: torch.device
-    data_source: MnistDataSource
-    model: Cnn
-
-    def __init__(self):
-        self.device = pytorch_gpu_utils.get_gpu_device()
-        self.data_source = MnistDataSource(Trainer.N_BATCH_SIZE)
-        self.model = Cnn().to(self.device)
-        # 展示模型结构。批次为指定批次数量，通道只有一个灰度通道，大小28x28。
-        torchinfo.summary(self.model, (Trainer.N_BATCH_SIZE, 1, 28, 28))
-
-    def train(self):
-        optimizer = torch.optim.Adam(self.model.parameters(), eps=1e-7)
-        # optimizer = torch.optim.AdamW(
-        #     self.model.parameters(),
-        #     lr=0.001,          # 两者默认学习率都是 0.001
-        #     betas=(0.9, 0.999), # 两者默认值相同
-        #     eps=1e-07,         # 【关键】匹配 TensorFlow 的默认 epsilon
-        #     weight_decay=0.0,   # 两者默认都是 0
-        #     amsgrad=False       # 两者默认都是 False
-        # )
-        loss_func = torch.nn.CrossEntropyLoss()
-
-        for epoch in range(Trainer.N_EPOCH):
-            self.model.train()
-
-            batch_images: torch.Tensor
-            batch_labels: torch.Tensor
-            for batch_index, (batch_images, batch_labels) in enumerate(self.data_source.train_loader):
-                gpu_images = batch_images.to(self.device)
-                gpu_labels = batch_labels.to(self.device)
-
-                prediction: torch.Tensor = self.model(gpu_images)
-                loss: torch.Tensor = loss_func(prediction, gpu_labels)
-                optimizer.zero_grad()
-                loss.backward()
-                optimizer.step()
-
-                if batch_index % 100 == 0:
-                    literal_loss = loss.item()
-                    print(f'Epoch: {epoch+1}, Batch: {batch_index}, Loss: {literal_loss:.4f}')
-
-    def save(self):
-        file_dir_path = Path(__file__).resolve().parent.parent / 'models'
-        file_dir_path.mkdir(parents=True, exist_ok=True)
-        file_path = file_dir_path / 'cnn.pth'
-        torch.save(self.model.state_dict(), file_path)
-        print(f'模型已保存至：{file_path}')
-
-    def test(self):
-        self.model.eval()
-        correct_sum = 0
-        total_sum = 0
-
-        with torch.no_grad():
-            for batch_images, batch_labels in self.data_source.test_loader:
-                gpu_images = batch_images.to(self.device)
-                gpu_labels = batch_labels.to(self.device)
-
-                possibilities: torch.Tensor = self.model(gpu_images)
-                # 输出出来是10个数字各自的可能性，所以要选取最高可能性的那个对比
-                # 在dim=1上找最大的那个，就选那个。dim=0是批次所以不管他。
-                _, prediction = possibilities.max(1)
-                # 返回标签的个数作为这一批的总个数
-                total_sum += gpu_labels.size(0)
-                correct_sum += prediction.eq(gpu_labels).sum()
-
-        test_acc = 100. * correct_sum / total_sum
-        print(f"准确率: {test_acc:.4f}%，共测试了{total_sum}张图片")
-
-def main():
-    trainer = Trainer()
-    trainer.train()
-    trainer.save()
-    trainer.test()
-
-# N_EPOCH: int = 5
-# N_BATCH_SIZE: int = 1000
-# N_LOG_INTERVAL: int = 10
-
 # class Trainer:
 
 #     device: torch.device
-#     data_source: MnistDataSource
+#     data_source: MnistDataLoaders
 #     model: Cnn
 
-#     trainer: Engine
-#     train_evaluator: Engine
-#     test_evaluator: Engine
-
 #     def __init__(self):
 #         self.device = gpu_utils.get_gpu_device()
+#         self.data_source = MnistDataLoaders(Trainer.N_BATCH_SIZE)
 #         self.model = Cnn().to(self.device)
-#         self.data_source = MnistDataSource(batch_size=N_BATCH_SIZE)
 #         # 展示模型结构。批次为指定批次数量，通道只有一个灰度通道，大小28x28。
-#         torchinfo.summary(self.model, (N_BATCH_SIZE, 1, 28, 28))
+#         torchinfo.summary(self.model, (Trainer.N_BATCH_SIZE, 1, 28, 28))
 
-#         #optimizer = torch.optim.Adam(self.model.parameters(), eps=1e-7)
-#         optimizer = torch.optim.AdamW(
-#             self.model.parameters(),
-#             lr=0.001,          # 两者默认学习率都是 0.001
-#             betas=(0.9, 0.999), # 两者默认值相同
-#             eps=1e-07,         # 【关键】匹配 TensorFlow 的默认 epsilon
-#             weight_decay=0.0,   # 两者默认都是 0
-#             amsgrad=False       # 两者默认都是 False
-#         )
-#         criterion = torch.nn.CrossEntropyLoss()
-#         self.trainer = ignite.engine.create_supervised_trainer(
-#             self.model, optimizer, criterion, self.device
-#         )
+#     def train(self):
+#         optimizer = torch.optim.Adam(self.model.parameters(), eps=1e-7)
+#         # optimizer = torch.optim.AdamW(
+#         #     self.model.parameters(),
+#         #     lr=0.001,          # 两者默认学习率都是 0.001
+#         #     betas=(0.9, 0.999), # 两者默认值相同
+#         #     eps=1e-07,         # 【关键】匹配 TensorFlow 的默认 epsilon
+#         #     weight_decay=0.0,   # 两者默认都是 0
+#         #     amsgrad=False       # 两者默认都是 False
+#         # )
+#         loss_func = torch.nn.CrossEntropyLoss()
 
-#         eval_metrics = {
-#             "accuracy": ignite.metrics.Accuracy(device=self.device),
-#             "loss": ignite.metrics.Loss(criterion, device=self.device)
-#         }
-#         self.train_evaluator = ignite.engine.create_supervised_evaluator(
-#             self.model, metrics=eval_metrics, device=self.device)
-#         self.test_evaluator = ignite.engine.create_supervised_evaluator(
-#             self.model, metrics=eval_metrics, device=self.device)
-        
-#         self.trainer.add_event_handler(
-#             Events.ITERATION_COMPLETED(every=N_LOG_INTERVAL),
-#             lambda engine: self.log_intrain_loss(engine)
-#         )
-#         self.trainer.add_event_handler(
-#             Events.EPOCH_COMPLETED,
-#             lambda trainer: self.log_train_results(trainer)
-#         )
-#         self.trainer.add_event_handler(
-#             Events.COMPLETED,
-#             lambda _: self.log_test_results()
-#         )
-#         self.trainer.add_event_handler(
-#             Events.COMPLETED,
-#             lambda _: self.save_model()
-#         )
+#         for epoch in range(Trainer.N_EPOCH):
+#             self.model.train()
 
-#         progressbar = ProgressBar()
-#         progressbar.attach(self.trainer)
+#             batch_images: torch.Tensor
+#             batch_labels: torch.Tensor
+#             for batch_index, (batch_images, batch_labels) in enumerate(self.data_source.train_loader):
+#                 gpu_images = batch_images.to(self.device)
+#                 gpu_labels = batch_labels.to(self.device)
 
-#     def run(self):
-#         self.trainer.run(self.data_source.train_loader, max_epochs=N_EPOCH)
+#                 prediction: torch.Tensor = self.model(gpu_images)
+#                 loss: torch.Tensor = loss_func(prediction, gpu_labels)
+#                 optimizer.zero_grad()
+#                 loss.backward()
+#                 optimizer.step()
 
-#     def log_intrain_loss(self, engine: Engine):
-#         print(f"Epoch: {engine.state.epoch}, Loss: {engine.state.output:.4f}\r", end="")
+#                 if batch_index % 100 == 0:
+#                     literal_loss = loss.item()
+#                     print(f'Epoch: {epoch+1}, Batch: {batch_index}, Loss: {literal_loss:.4f}')
 
-#     def log_train_results(self, trainer: Engine):
-#         self.train_evaluator.run(self.data_source.train_loader)
-#         metrics = self.train_evaluator.state.metrics
-#         print()
-#         print(f"Training - Epoch: {trainer.state.epoch}, Avg Accuracy: {metrics['accuracy']:.4f}, Avg Loss: {metrics['loss']:.4f}")
-
-#     def log_test_results(self):
-#         self.test_evaluator.run(self.data_source.test_loader)
-#         metrics = self.test_evaluator.state.metrics
-#         print(f"Test - Avg Accuracy: {metrics['accuracy']:.4f} Avg Loss: {metrics['loss']:.4f}")
-
-#     def save_model(self):
+#     def save(self):
 #         file_dir_path = Path(__file__).resolve().parent.parent / 'models'
 #         file_dir_path.mkdir(parents=True, exist_ok=True)
 #         file_path = file_dir_path / 'cnn.pth'
 #         torch.save(self.model.state_dict(), file_path)
-#         print(f'Model was saved into: {file_path}')
+#         print(f'模型已保存至：{file_path}')
+
+#     def test(self):
+#         self.model.eval()
+#         correct_sum = 0
+#         total_sum = 0
+
+#         with torch.no_grad():
+#             batch_images: torch.Tensor
+#             batch_labels: torch.Tensor
+#             for batch_images, batch_labels in self.data_source.test_loader:
+#                 gpu_images = batch_images.to(self.device)
+#                 gpu_labels = batch_labels.to(self.device)
+
+#                 possibilities: torch.Tensor = self.model(gpu_images)
+#                 # 输出出来是10个数字各自的可能性，所以要选取最高可能性的那个对比
+#                 # 在dim=1上找最大的那个，就选那个。dim=0是批次所以不管他。
+#                 _, prediction = possibilities.max(1)
+#                 # 返回标签的个数作为这一批的总个数
+#                 total_sum += gpu_labels.size(0)
+#                 correct_sum += prediction.eq(gpu_labels).sum()
+
+#         test_acc = 100. * correct_sum / total_sum
+#         print(f"准确率: {test_acc:.4f}%，共测试了{total_sum}张图片")
 
 # def main():
 #     trainer = Trainer()
-#     trainer.run()
+#     trainer.train()
+#     trainer.save()
+#     trainer.test()
 
+class Trainer:
+    """核心训练器"""
+
+    device: torch.device
+    data_source: MnistDataLoaders
+    model: Cnn
+
+    trainer: Engine
+    trainer_accuracy: Accuracy
+    evaluator: Engine
+    pbar: ProgressBar
+
+    def __init__(self):
+        # 创建训练设备，模型和数据加载器。
+        self.device = gpu_utils.get_gpu_device()
+        self.model = Cnn().to(self.device)
+        self.data_source = MnistDataLoaders(batch_size=settings.N_BATCH_SIZE)
+        # 展示模型结构。批次为指定批次数量，通道只有一个灰度通道，大小28x28。
+        torchinfo.summary(self.model, (settings.N_BATCH_SIZE, 1, 28, 28))
+
+        # 优化器和损失函数
+        optimizer = torch.optim.Adam(self.model.parameters(), eps=1e-7)
+        criterion = torch.nn.CrossEntropyLoss()
+        # 创建训练器
+        self.trainer = ignite.engine.create_supervised_trainer(
+            self.model, optimizer, criterion, self.device,
+            # 输出转换为这种形式，因为后面的测量需要用其中一些参数。
+            # 默认的输出只输出loss。
+            output_transform=lambda x, y, y_pred, loss: {
+                'loss': loss.item(),
+                'y': y,
+                'y_pred': y_pred
+            }
+        )
+        # 设置训练器测量数据
+        self.trainer_accuracy = Accuracy(
+            device=self.device,
+            # 转换为Accuracy需要的形式。
+            output_transform=lambda o: (o['y_pred'], o['y'])
+        )
+        self.trainer_accuracy.attach(self.trainer, 'accuracy')
+        # YYC MARK: 这里要手动reset一下，不然第一次运行没有accuracy
+        self.trainer_accuracy.reset()
+        # 每次epoch前重置accuracy
+        self.trainer.add_event_handler(
+            Events.EPOCH_STARTED,
+            lambda: self.trainer_accuracy.reset()
+        )
+        # 将训练器关联到进度条
+        self.pbar = ProgressBar(persist=True)
+        self.pbar.attach(self.trainer, 
+                         metric_names=['accuracy'],
+                         output_transform=lambda o: {"loss": o['loss']})
+        # 训练完毕后保存模型
+        self.trainer.add_event_handler(
+            Events.COMPLETED,
+            lambda: self.save_model()
+        )
+
+        # 创建测试的评估器的评估量
+        evaluator_metrics = {
+            "accuracy": ignite.metrics.Accuracy(device=self.device),
+            "loss": ignite.metrics.Loss(criterion, device=self.device)
+        }
+        # 创建测试评估器
+        self.evaluator = ignite.engine.create_supervised_evaluator(
+            self.model, metrics=evaluator_metrics, device=self.device)
+        
+    def run(self):
+        # 训练模型
+        self.trainer.run(self.data_source.train_loader, max_epochs=settings.N_EPOCH)
+        # 测试模型并输出结果
+        self.evaluator.run(self.data_source.test_loader)
+        metrics = self.evaluator.state.metrics
+        print(f"Test Done. Accuracy: {metrics['accuracy']:.4f} Loss: {metrics['loss']:.4f}")
+
+    def save_model(self):
+        # 确保保存模型的文件夹存在。
+        settings.SAVED_MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
+        # 仅保存模型参数
+        torch.save(self.model.state_dict(), settings.SAVED_MODEL_PATH)
+        print(f'Model was saved into: {settings.SAVED_MODEL_PATH}')
+
+def main():
+    trainer = Trainer()
+    trainer.run()
 
 
 if __name__ == "__main__":
-    pytorch_gpu_utils.print_gpu_availability()
+    gpu_utils.print_gpu_availability()
     main()