yyc12345/ai-school
1
0
Fork 0
Code Activity

update for change of exp2 and add exp3

Browse Source
This commit is contained in:
yyc12345
2025-11-30 16:24:32 +08:00
parent af890d899e
commit 48fcdfcc80
17 changed files with 859 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

240
exp2/modified/train.py
View File

@@ -1,145 +1,57 @@
from pathlib import Path
import sys
import typing
import numpy
import torch
from torch.utils.data import DataLoader, Dataset
from torchvision.transforms import v2 as tvtrans
import matplotlib.pyplot as plt
import torch.nn.functional as F
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 CNN(torch.nn.Module):
"""卷积神经网络模型"""
def __init__(self):
super(CNN, self).__init__()
self.conv1 = torch.nn.Conv2d(1, 32, kernel_size=(3, 3))
self.pool1 = torch.nn.MaxPool2d(kernel_size=(2, 2))
self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=(3, 3))
self.pool2 = torch.nn.MaxPool2d(kernel_size=(2, 2))
self.conv3 = torch.nn.Conv2d(64, 64, kernel_size=(3, 3))
self.flatten = torch.nn.Flatten()
# 28x28过第一轮卷积后变为26x26，过第一轮池化后变为13x13
# 过第二轮卷积后变为11x11，过第二轮池化后变为5x5
# 过第三轮卷积后变为3x3。
# 最后一轮卷积核个数为64。
self.fc1 = torch.nn.Linear(64 * 3 * 3, 64)
self.fc2 = torch.nn.Linear(64, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool1(x)
x = F.relu(self.conv2(x))
x = self.pool2(x)
x = F.relu(self.conv3(x))
x = self.flatten(x)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.softmax(x, dim=1)
class MnistDataset(Dataset):
"""用于加载Mnist的自定义数据集"""
shape: int
transform: tvtrans.Transform
images_data: numpy.ndarray
labels_data: torch.Tensor
def __init__(self, images: numpy.ndarray, labels: numpy.ndarray, transform: tvtrans.Transform):
images_len: int = images.size(0)
labels_len: int = labels.size(0)
assert (images_len == labels_len)
self.shape = images_len
self.images_data = images
self.labels_data = torch.from_numpy(labels)
self.transform = transform
def __getitem__(self, index):
return self.transform(self.images_data[index]), self.labels_data[index]
def __len__(self):
return self.shape
class DataSource:
"""用于读取MNIST数据的数据读取器"""
train_data: DataLoader
test_data: DataLoader
def __init__(self, batch_size: int):
datasets_path = Path(__file__).resolve().parent.parent / 'datasets' / 'mnist.npz'
datasets = numpy.load(datasets_path)
# 所有图片均为黑底白字
# 6万张训练图片：60000x28x28。标签只有第一维。
train_images = datasets['x_train']
train_labels = datasets['y_train']
# 1万张测试图片：10000x28x28。标签只有第一维。
test_images = datasets['x_test']
test_labels = datasets['y_test']
# 定义数据转换器
trans = tvtrans.Compose([
# 从uint8转换为float32并自动归一化到0-1区间
# tvtrans.ToTensor(),
tvtrans.ToImage(),
tvtrans.ToDtype(torch.float32, scale=True),
# 为了符合后面图像的输入颜色通道条件，要在最后挤出一个新的维度
#tvtrans.Lambda(lambda x: x.unsqueeze(-1))
])
# 创建数据集
train_dataset = MnistDataset(train_images,
train_labels,
transform=trans)
test_dataset = MnistDataset(test_images, test_labels, transform=trans)
# 赋值到自身
self.train_data = DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=False)
self.test_data = DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)
class Trainer:
N_EPOCH: typing.ClassVar[int] = 5
N_BATCH_SIZE: typing.ClassVar[int] = 1000
device: torch.device
data_source: DataSource
cnn: CNN
data_source: MnistDataSource
model: CNN
def __init__(self):
self.device = gpu_utils.get_gpu_device()
self.data_source = DataSource(Trainer.N_BATCH_SIZE)
self.cnn = CNN().to(self.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.cnn.parameters())
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.cnn.train()
self.model.train()
batch_images: torch.Tensor
batch_labels: torch.Tensor
for batch_index, (batch_images, batch_labels) in enumerate(self.data_source.train_data):
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)
optimizer.zero_grad()
prediction: torch.Tensor = self.cnn(gpu_images)
prediction: torch.Tensor = self.model(gpu_images)
loss: torch.Tensor = loss_func(prediction, gpu_labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
@@ -151,20 +63,20 @@ class Trainer:
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.cnn.state_dict(), file_path)
torch.save(self.model.state_dict(), file_path)
print(f'模型已保存至：{file_path}')
def test(self):
self.cnn.eval()
self.model.eval()
correct_sum = 0
total_sum = 0
with torch.no_grad():
for batch_images, batch_labels in self.data_source.test_data:
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.cnn(gpu_images)
possibilities: torch.Tensor = self.model(gpu_images)
# 输出出来是10个数字各自的可能性，所以要选取最高可能性的那个对比
# 在dim=1上找最大的那个，就选那个。dim=0是批次所以不管他。
_, prediction = possibilities.max(1)
@@ -175,13 +87,105 @@ class Trainer:
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()