Upload train.py

train.py

@@ -0,0 +1,113 @@
+import torch
+import torchvision
+from torch import optim
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from Pytorch_MNIST图片识别.model import Net
+import matplotlib.pyplot as plt
+import os
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+# TODO epoch的数量定义了我们将循环整个训练数据集的次数
+n_epochs = 3
+
+# TODO 使用batch_size=64进行训练，并使用size=1000对这个数据集进行测试
+batch_size_train = 64
+batch_size_test = 1000
+
+# TODO 优化器的超参数
+learning_rate = 0.01
+momentum = 0.5
+
+log_interval = 10
+random_seed = 1
+torch.manual_seed(random_seed)
+
+# TODO 自动将MNIST数据集下载到目录下的data文件夹
+train_loader = torch.utils.data.DataLoader(
+    torchvision.datasets.MNIST('./data/', train=True, download=True,
+                               transform=torchvision.transforms.Compose([
+
+                                   torchvision.transforms.ToTensor(),
+                                   # TODO MNIST数据集的全局平均值和标准偏差
+                                   torchvision.transforms.Normalize(
+                                       (0.1307,), (0.3081,))
+                               ])),
+    batch_size=batch_size_train, shuffle=True)
+test_loader = torch.utils.data.DataLoader(
+    torchvision.datasets.MNIST('./data/', train=False, download=True,
+                               transform=torchvision.transforms.Compose([
+                                   torchvision.transforms.ToTensor(),
+                                   # TODO MNIST数据集的全局平均值和标准偏差
+                                   torchvision.transforms.Normalize(
+                                       (0.1307,), (0.3081,))
+                               ])),
+    batch_size=batch_size_test, shuffle=False)
+
+# TODO 初始化网络和优化器
+network = Net()
+optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)
+
+train_losses = []
+train_counter = []
+test_losses = []
+test_counter = [i * len(train_loader.dataset) for i in range(n_epochs + 1)]
+
+# TODO 模型存储位置(一个是完整的模型，一个是只有参数的模型）
+# TODO 需要先建立一个model文件夹
+model_path = './model1/model.pth'
+optimizer_path = './model1/optimizer.pth'
+
+
+def train(epoch):
+    network.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        # TODO 需要使用optimizer.zero_grad()手动将梯度设置为零，因为PyTorch在默认情况下会累积梯度
+        optimizer.zero_grad()
+        output = network(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % log_interval == 0:
+            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item()))
+            train_losses.append(loss.item())
+            train_counter.append(
+                (batch_idx * 64) + ((epoch - 1) * len(train_loader.dataset)))
+    if epoch == (n_epochs - 1):
+        # TODO 存储模型
+        torch.save(network.state_dict(), model_path)
+        torch.save(optimizer.state_dict(), optimizer_path)
+
+
+def test():
+    network.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            output = network(data)
+            test_loss += F.nll_loss(output, target, size_average=False).item()
+            pred = output.data.max(1, keepdim=True)[1]
+            correct += pred.eq(target.data.view_as(pred)).sum()
+    test_loss /= len(test_loader.dataset)
+    test_losses.append(test_loss)
+    print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
+        test_loss, correct, len(test_loader.dataset),
+        100. * correct / len(test_loader.dataset)))
+
+
+for epoch in range(1, n_epochs + 1):
+    train(epoch)
+    test()
+
+fig = plt.figure()
+plt.plot(train_counter, train_losses, color='blue')
+print(len(test_counter))
+print(len(test_losses))
+plt.scatter(test_counter, test_losses, color='red')
+plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
+plt.xlabel('number of training examples seen')
+plt.ylabel('negative log likelihood loss')
+plt.show()