|
@@ -0,0 +1,93 @@
|
|
|
+import torch
|
|
|
+import torch.nn as nn
|
|
|
+import torch.optim as optim
|
|
|
+from torchvision import datasets, transforms
|
|
|
+import matplotlib.pyplot as plt
|
|
|
+
|
|
|
+# 检查是否有可用的GPU
|
|
|
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
+
|
|
|
+# 定义数据预处理步骤
|
|
|
+transform = transforms.Compose([
|
|
|
+ transforms.ToTensor(),
|
|
|
+ transforms.Normalize((0.1307,), (0.3081,))
|
|
|
+])
|
|
|
+
|
|
|
+# 加载MNIST训练数据集
|
|
|
+train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
|
|
|
+train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
|
|
|
+
|
|
|
+# 加载MNIST测试数据集
|
|
|
+test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)
|
|
|
+test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)
|
|
|
+
|
|
|
+# 定义卷积神经网络模型
|
|
|
+class Net(nn.Module):
|
|
|
+ def __init__(self):
|
|
|
+ super(Net, self).__init__()
|
|
|
+ self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
|
|
|
+ self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
|
|
|
+ self.fc1 = nn.Linear(320, 50)
|
|
|
+ self.fc2 = nn.Linear(50, 10)
|
|
|
+
|
|
|
+ def forward(self, x):
|
|
|
+ x = nn.functional.relu(nn.functional.max_pool2d(self.conv1(x), 2))
|
|
|
+ x = nn.functional.relu(nn.functional.max_pool2d(self.conv2(x), 2))
|
|
|
+ x = x.view(-1, 320)
|
|
|
+ x = nn.functional.relu(self.fc1(x))
|
|
|
+ x = self.fc2(x)
|
|
|
+ return x
|
|
|
+
|
|
|
+# 创建模型实例并将其移动到设备(GPU或CPU)
|
|
|
+model = Net()
|
|
|
+model.to(device)
|
|
|
+
|
|
|
+# 定义损失函数和优化器
|
|
|
+criterion = nn.CrossEntropyLoss()
|
|
|
+optimizer = optim.SGD(model.parameters(), lr=0.01)
|
|
|
+
|
|
|
+# 训练模型
|
|
|
+def train_model():
|
|
|
+ model.train()
|
|
|
+ for epoch in range(20):
|
|
|
+ running_loss = 0.0
|
|
|
+ for batch_idx, (data, target) in enumerate(train_loader):
|
|
|
+ data, target = data.to(device), target.to(device)
|
|
|
+
|
|
|
+ optimizer.zero_grad()
|
|
|
+
|
|
|
+ output = model(data)
|
|
|
+ loss = criterion(output, target)
|
|
|
+ loss.backward()
|
|
|
+ optimizer.step()
|
|
|
+
|
|
|
+ running_loss += loss.item()
|
|
|
+ if batch_idx % 100 == 0:
|
|
|
+ print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
|
|
|
+ epoch + 1, batch_idx * 64, len(train_loader.dataset),
|
|
|
+ 100. * batch_idx / len(train_loader), running_loss / (batch_idx + 1)))
|
|
|
+
|
|
|
+ print('====> Epoch: {} Average Loss: {:.4f}'.format(epoch + 1, running_loss / len(train_loader)))
|
|
|
+
|
|
|
+# 测试模型
|
|
|
+def test_model():
|
|
|
+ model.eval()
|
|
|
+ test_loss = 0
|
|
|
+ correct = 0
|
|
|
+ with torch.no_grad():
|
|
|
+ for data, target in test_loader:
|
|
|
+ data, target = data.to(device), target.to(device)
|
|
|
+
|
|
|
+ output = model(data)
|
|
|
+ test_loss += criterion(output, target).item()
|
|
|
+ pred = output.argmax(dim=1, keepdim=True)
|
|
|
+ correct += pred.eq(target.view_as(pred)).sum().item()
|
|
|
+
|
|
|
+ test_loss /= len(test_loader)
|
|
|
+ accuracy = correct / len(test_loader.dataset)
|
|
|
+
|
|
|
+ print('Test set: Average loss: {:.4f}, Accuracy: {:.2f}%'.format(test_loss, 100 * accuracy))
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ train_model()
|
|
|
+ test_model()
|