代码:
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.nn.functional as Fbatch_size = 64
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) #把原始图像转为tensor 这是均值和方差train_set = datasets.MNIST(root='./data/mnist', train=True, download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)test_set = datasets.MNIST(root='./data/mnist', train=False, download=True, transform=transform)
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=True)class LinearClassifier(nn.Module):def __init__(self):super(LinearClassifier, self).__init__()self.l1 = nn.Linear(784, 512)self.l2 = nn.Linear(512, 256)self.l3 = nn.Linear(256, 128)self.l4 = nn.Linear(128, 64)self.l5 = nn.Linear(64, 10)def forward(self, x):x = x.view(-1, 784)x = F.relu(self.l1(x))x = F.relu(self.l2(x))x = F.relu(self.l3(x))x = F.relu(self.l4(x))return self.l5(x)criterion = nn.CrossEntropyLoss()
model = LinearClassifier()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)def train(epoch):running_loss = 0.0for i, data in enumerate(train_loader, 0):inputs, labels = dataoptimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()if i % 300 == 299: # print every 300 mini-batchesprint('[%d, %5d] loss: %.3f' %(epoch + 1, i + 1, running_loss / 300))running_loss = 0.0def test(epoch):correct = 0total = 0with torch.no_grad():for data in test_loader:images, labels = dataoutputs = model(images)_, predicted = torch.max(outputs.data, 1)total += labels.size(0)correct += (predicted == labels).sum().item()print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))if __name__ == '__main__':for epoch in range(100):train(epoch)if epoch % 10 == 0:test(epoch)