117 lines
3.5 KiB
Python
117 lines
3.5 KiB
Python
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()
|