diff --git a/Lab-2/src/data_transform.py b/Lab-2/src/data_transform.py index 8a7eeb7..b81ba26 100644 --- a/Lab-2/src/data_transform.py +++ b/Lab-2/src/data_transform.py @@ -18,7 +18,7 @@ def get_transform(): ]) -def load_mnist(data_size=100, batch_size=16, seed=42, root='../data'): +def load_mnist(data_size=100, batch_size=16, seed=42, root='../data', train=True): """ 加载 MNIST 数据集。 @@ -27,32 +27,33 @@ def load_mnist(data_size=100, batch_size=16, seed=42, root='../data'): - batch_size (int): 每批次处理的数据数量。 - seed (int): 随机种子，用于保证数据切割的可重复性。 - root (str): 数据集保存的路径。 + - train (bool): 是否加载训练集。 Returns: - - train_loader (DataLoader): 训练数据集加载器。 + - DataLoader: 数据加载器。 """ transform = get_transform() - train_dataset = torchvision.datasets.MNIST( + dataset = torchvision.datasets.MNIST( root=root, - train=True, + train=train, download=True, transform=transform ) - data_train, _ = random_split( - dataset=train_dataset, - lengths=[data_size, len(train_dataset) - data_size], + data_subset, _ = random_split( + dataset=dataset, + lengths=[data_size, len(dataset) - data_size], generator=torch.Generator().manual_seed(seed) ) - train_loader = DataLoader( - dataset=data_train, + data_loader = DataLoader( + dataset=data_subset, batch_size=batch_size, shuffle=True ) - return train_loader + return data_loader def load_and_preview_mnist(data_size=100, batch_size=16, seed=42, root='../data'): @@ -65,14 +66,14 @@ def load_and_preview_mnist(data_size=100, batch_size=16, seed=42, root='../data' - seed (int): 随机种子，用于保证数据切割的可重复性。 - root (str): 数据集保存的路径。 """ - train_loader = load_mnist(data_size=data_size, batch_size=batch_size, seed=seed, root=root) + train_loader = load_mnist(data_size=data_size, batch_size=batch_size, seed=seed, root=root, train=True) images, labels = next(iter(train_loader)) img_grid = torchvision.utils.make_grid(images, nrow=batch_size // 4) img = img_grid.numpy().transpose(1, 2, 0) - mean = [0.5, 0.5, 0.5] - std = [0.5, 0.5, 0.5] + mean = [0.5] + std = [0.5] img = img * std + mean print([labels[i] for i in range(batch_size)]) @@ -85,4 +86,4 @@ if __name__ == '__main__': load_and_preview_mnist() # Or use custom parameters: - # load_and_preview_mnist(data_size=200, batch_size=8, seed=123, root='../data') + # load_and_preview_mnist(data_size=200, batch_size=8, seed=123, root='./data') diff --git a/Lab-2/src/model_training.py b/Lab-2/src/model_training.py new file mode 100644 index 0000000..696bf56 --- /dev/null +++ b/Lab-2/src/model_training.py @@ -0,0 +1,116 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from data_transform import load_mnist, load_and_preview_mnist +import matplotlib.pyplot as plt +import torchvision.utils as vutils + + +# 定义卷积神经网络模型 +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.conv1 = nn.Sequential( + nn.Conv2d(1, 64, kernel_size=3, stride=2, padding=1), + nn.ReLU(), + nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), + nn.ReLU(), + # nn.MaxPool2d(stride=2, kernel_size=2) + ) + + self.dense = nn.Sequential( + nn.Linear(7 * 7 * 128, 512), + nn.ReLU(), + nn.Dropout(p=0.8), + nn.Linear(512, 10) + ) + + def forward(self, x): + x = self.conv1(x) # 卷积处理 + x = x.view(-1, 7 * 7 * 128) # 对参数实行扁平化处理 + x = self.dense(x) + return x + + +# 加载和预览数据集 +def main(): + # 参数设置 + data_size = 1000 + batch_size = 4 + seed = 42 + root = '../data' + epochs_n = 5 + learning_rate = 0.001 + + # 加载数据加载器 + train_loader = load_mnist(data_size=data_size, batch_size=batch_size, seed=seed, root=root, train=True) + test_loader = load_mnist(data_size=data_size, batch_size=batch_size, seed=seed, root=root, train=False) + + # 预览数据集（可选） + load_and_preview_mnist(data_size=data_size, batch_size=batch_size, seed=seed, root=root) + + # 初始化模型 + model = Model() + cost = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=learning_rate) + + # 训练模型 + for epoch in range(epochs_n): + model.train() + running_loss = 0.0 + running_correct = 0 + print(f"Epoch {epoch + 1}/{epochs_n}") + print("-" * 10) + + for data in train_loader: + images, labels = data + optimizer.zero_grad() + outputs = model(images) + _, preds = torch.max(outputs.data, 1) + loss = cost(outputs, labels) + loss.backward() + optimizer.step() + running_loss += loss.item() + running_correct += (preds == labels).sum().item() + + epoch_loss = running_loss / len(train_loader.dataset) + epoch_acc = 100 * running_correct / len(train_loader.dataset) + + print(f"Loss: {epoch_loss:.4f}, Train Accuracy: {epoch_acc:.4f}%") + + # 测试模型 + model.eval() + testing_correct = 0 + with torch.no_grad(): + for data in test_loader: + images, labels = data + outputs = model(images) + _, preds = torch.max(outputs.data, 1) + testing_correct += (preds == labels).sum().item() + + test_acc = 100 * testing_correct / len(test_loader.dataset) + print(f"Test Accuracy: {test_acc:.4f}%") + + # 测试数据可视化 + X_test, y_test = next(iter(test_loader)) + with torch.no_grad(): + pred = model(X_test) + _, predicted_labels = torch.max(pred, 1) + + print("Predicted Labels:", [i.item() for i in predicted_labels]) + print("True Labels:", [i.item() for i in y_test]) + + # 显示图像 + img_grid = vutils.make_grid(X_test, nrow=batch_size // 4) + img = img_grid.numpy().transpose(1, 2, 0) + mean = [0.5] + std = [0.5] + img = img * std + mean + plt.imshow(img, cmap='gray') + plt.axis('off') + plt.show() + + +if __name__ == '__main__': + main()