use ignite for exp2
This commit is contained in:
@@ -3,9 +3,10 @@ import numpy
|
||||
import torch
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torchvision.transforms import v2 as tvtrans
|
||||
import settings
|
||||
|
||||
class MnistDataset(Dataset):
|
||||
"""用于加载Mnist数据的自定义数据集"""
|
||||
"""适配PyTorch的自定义Dataset，用于加载MNIST数据。"""
|
||||
|
||||
shape: int
|
||||
transform: tvtrans.Transform
|
||||
@@ -29,15 +30,14 @@ class MnistDataset(Dataset):
|
||||
return self.shape
|
||||
|
||||
|
||||
class MnistDataSource:
|
||||
"""用于读取MNIST数据的数据读取器"""
|
||||
class MnistDataLoaders:
|
||||
"""包含适配PyTorch的训练数据Loader和测试数据Loader的类。"""
|
||||
|
||||
train_loader: DataLoader
|
||||
test_loader: DataLoader
|
||||
|
||||
def __init__(self, batch_size: int):
|
||||
dataset_path = Path(__file__).resolve().parent.parent / 'datasets' / 'mnist.npz'
|
||||
dataset = numpy.load(dataset_path)
|
||||
dataset = numpy.load(settings.MNIST_DATASET_PATH)
|
||||
|
||||
# 所有图片均为黑底白字
|
||||
# 6万张训练图片：60000x28x28。标签只有第一维。
|
||||
@@ -50,16 +50,19 @@ class MnistDataSource:
|
||||
# 定义数据转换器
|
||||
trans = tvtrans.Compose([
|
||||
# 从uint8转换为float32并自动归一化到0-1区间
|
||||
# YYC MARK: 下面这个被标outdated了，换下面两个替代。
|
||||
# tvtrans.ToTensor(),
|
||||
tvtrans.ToImage(),
|
||||
tvtrans.ToDtype(torch.float32, scale=True),
|
||||
|
||||
# 为了符合后面图像的输入颜色通道条件，要在最后挤出一个新的维度
|
||||
#tvtrans.Lambda(lambda x: x.unsqueeze(-1))
|
||||
# YYC MARK: 上面这两步已经帮我们自动挤出那个灰度通道了。
|
||||
# tvtrans.Lambda(lambda x: x.unsqueeze(-1))
|
||||
|
||||
# 这个特定的标准化参数 (0.1307, 0.3081) 是 MNIST 数据集的标准化参数，这些数值是MNIST训练集的全局均值和标准差。
|
||||
# 这种标准化有助于模型训练时的数值稳定性和收敛速度。
|
||||
#tvtrans.Normalize((0.1307,), (0.3081,)),
|
||||
# YYC MARK: 但我不想用，反正最后训练的也收敛。
|
||||
# tvtrans.Normalize((0.1307,), (0.3081,)),
|
||||
])
|
||||
|
||||
# 创建数据集
|
||||
@@ -69,9 +72,9 @@ class MnistDataSource:
|
||||
transform=trans)
|
||||
|
||||
# 赋值到自身
|
||||
self.train_loader = DataLoader(dataset=train_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False)
|
||||
self.test_loader = DataLoader(dataset=test_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False)
|
||||
self.train_loader = DataLoader(dataset=train_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False)
|
||||
self.test_loader = DataLoader(dataset=test_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False)
|
||||
|
||||
@@ -21,9 +21,9 @@ class Cnn(torch.nn.Module):
|
||||
self.fc2 = torch.nn.Linear(64, 10)
|
||||
|
||||
# 初始化模型参数
|
||||
self._initialize_weights()
|
||||
self.__initialize_weights()
|
||||
|
||||
def _initialize_weights(self):
|
||||
def __initialize_weights(self):
|
||||
# YYC MARK:
|
||||
# 把两个全连接线性层按tensorflow默认设置初始化，即：
|
||||
# - kernel_initializer='glorot_uniform'
|
||||
|
||||
@@ -8,26 +8,36 @@ import matplotlib.pyplot as plt
|
||||
from model import Cnn
|
||||
|
||||
sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
|
||||
import pytorch_gpu_utils
|
||||
import gpu_utils
|
||||
|
||||
|
||||
class PredictResult:
|
||||
"""预测的结果"""
|
||||
|
||||
possibilities: torch.Tensor
|
||||
"""预测结果，是每个数字不同的概率，是经过softmax后的数值"""
|
||||
"""每个数字不同的概率"""
|
||||
|
||||
def __init__(self, possibilities: torch.Tensor):
|
||||
"""
|
||||
创建预测结果。
|
||||
|
||||
:param possibilities: 传入的tensor表示每个数字不同的概率，是经过softmax后的数值。
|
||||
其shape为二维。dim 0为batch，应当只有一维；dim 1为每个数字对应的概率。
|
||||
"""
|
||||
self.possibilities = possibilities
|
||||
|
||||
def chosen_number(self) -> int:
|
||||
"""获取最终选定的数字"""
|
||||
# 依然是找最大的那个index
|
||||
_, prediction = self.possibilities.max(1)
|
||||
return prediction.item()
|
||||
# 输出出来是10个数字各自的可能性，所以要选取最高可能性的那个对比
|
||||
# 在dim=1上找最大的那个，就选那个。dim=0是批次所以不管他。
|
||||
return self.possibilities.argmax(1).item()
|
||||
|
||||
def number_possibilities(self) -> list[float]:
|
||||
"""获取每个数字出现的概率"""
|
||||
"""
|
||||
获取每个数字出现的概率
|
||||
|
||||
:return: 返回一个具有10个元素的列表，列表的每一项表示当前index所代表数字的概率。
|
||||
"""
|
||||
return list(self.possibilities[0][i].item() for i in range(10))
|
||||
|
||||
class Predictor:
|
||||
@@ -35,14 +45,14 @@ class Predictor:
|
||||
model: Cnn
|
||||
|
||||
def __init__(self):
|
||||
self.device = pytorch_gpu_utils.get_gpu_device()
|
||||
self.device = gpu_utils.get_gpu_device()
|
||||
self.model = Cnn().to(self.device)
|
||||
|
||||
# 加载保存好的模型参数
|
||||
file_path = Path(__file__).resolve().parent.parent / 'models' / 'cnn.pth'
|
||||
self.model.load_state_dict(torch.load(file_path))
|
||||
|
||||
def generic_predict(self, in_data: torch.Tensor) -> PredictResult:
|
||||
def __predict(self, in_data: torch.Tensor) -> PredictResult:
|
||||
"""
|
||||
其它预测函数都要使用的预测后端。其它预测函数将数据处理成Tensor，然后传递给此函数进行实际预测。
|
||||
|
||||
@@ -62,24 +72,36 @@ class Predictor:
|
||||
|
||||
|
||||
def predict_sketchpad(self, image: list[list[bool]]) -> PredictResult:
|
||||
"""
|
||||
以sketchpad的数据进行预测。
|
||||
|
||||
:param image: 该列表的shape必须为28x28
|
||||
"""
|
||||
input = torch.Tensor(image).float()
|
||||
assert(input.dim() == 2)
|
||||
assert(input.size(0) == 28)
|
||||
assert(input.size(1) == 28)
|
||||
|
||||
return self.generic_predict(input)
|
||||
return self.__predict(input)
|
||||
|
||||
def predict_image(self, image: ImageFile.ImageFile) -> PredictResult:
|
||||
# 确保图像为灰度图像，然后转换为numpy数组。
|
||||
# 注意这里的numpy数组是只读的，所以要先拷贝一份
|
||||
"""
|
||||
以Pillow图像的数据进行预测。
|
||||
|
||||
:param image: Pillow图像数据。该图像必须为28x28大小。
|
||||
"""
|
||||
# 确保图像为灰度图像，以及宽高合适
|
||||
grayscale_image = image.convert('L')
|
||||
assert(grayscale_image.width == 28)
|
||||
assert(grayscale_image.height == 28)
|
||||
# 转换为numpy数组。注意这里的numpy数组是只读的，所以要先拷贝一份
|
||||
numpy_data = numpy.reshape(grayscale_image, (28, 28), copy=True)
|
||||
# 转换到Tensor，设置dtype
|
||||
data = torch.from_numpy(numpy_data).float()
|
||||
# 归一化到255，又因为图像输入是白底黑字，需要做转换。
|
||||
data.div_(255.0).sub_(1).mul_(-1)
|
||||
|
||||
return self.generic_predict(data)
|
||||
return self.__predict(data)
|
||||
|
||||
def main():
|
||||
predictor = Predictor()
|
||||
@@ -101,6 +123,6 @@ def main():
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytorch_gpu_utils.print_gpu_availability()
|
||||
gpu_utils.print_gpu_availability()
|
||||
main()
|
||||
|
||||
|
||||
12
exp2/modified/settings.py
Normal file
12
exp2/modified/settings.py
Normal file
@@ -0,0 +1,12 @@
|
||||
from pathlib import Path
|
||||
|
||||
MNIST_DATASET_PATH: Path = Path(__file__).resolve().parent.parent / 'datasets' / 'mnist.npz'
|
||||
"""MNIST数据集文件的路径"""
|
||||
|
||||
SAVED_MODEL_PATH: Path = Path(__file__).resolve().parent.parent / 'models' / 'cnn.pth'
|
||||
"""训练好的模型保存的位置"""
|
||||
|
||||
N_EPOCH: int = 5
|
||||
"""训练时的epoch次数"""
|
||||
N_BATCH_SIZE: int = 1000
|
||||
"""训练时的batch size"""
|
||||
@@ -5,7 +5,7 @@ import tkinter as tk
|
||||
from predict import PredictResult, Predictor
|
||||
|
||||
sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
|
||||
import pytorch_gpu_utils
|
||||
import gpu_utils
|
||||
|
||||
|
||||
class SketchpadApp:
|
||||
@@ -181,7 +181,7 @@ class SketchpadApp:
|
||||
# endregion
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytorch_gpu_utils.print_gpu_availability()
|
||||
gpu_utils.print_gpu_availability()
|
||||
|
||||
predictor = Predictor()
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user