193 lines
7.2 KiB
Python
193 lines
7.2 KiB
Python
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 mnist import CNN, MnistDataSource
|
||
|
||
sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
|
||
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 = 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__":
|
||
gpu_utils.print_gpu_availability()
|
||
main()
|