ai-school/exp2/modified/train.py

from pathlib import Path
import sys
import typing
import torch
import torchinfo
import ignite.engine
import ignite.metrics
from ignite.engine import Engine, Events
from ignite.handlers.tqdm_logger import ProgressBar
from dataset import MnistDataSource
from model import Cnn

sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
import pytorch_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
#     model: Cnn

#     trainer: Engine
#     train_evaluator: Engine
#     test_evaluator: Engine

#     def __init__(self):
#         self.device = gpu_utils.get_gpu_device()
#         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))

#         #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
#         )

#         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()
#         )

#         progressbar = ProgressBar()
#         progressbar.attach(self.trainer)

#     def run(self):
#         self.trainer.run(self.data_source.train_loader, max_epochs=N_EPOCH)

#     def log_intrain_loss(self, engine: Engine):
#         print(f"Epoch: {engine.state.epoch}, Loss: {engine.state.output:.4f}\r", end="")

#     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):
#         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}')

# def main():
#     trainer = Trainer()
#     trainer.run()



if __name__ == "__main__":
    pytorch_gpu_utils.print_gpu_availability()
    main()