# 导入相关库函数
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

# 设备配置
device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 超参数
input_size = 784
hidden_size = 500
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 0.001

# MNIST数据集
train_dataset = torchvision.datasets.MNIST(root='../../data',
                                         train=True,
                                         transform=transforms.ToTensor(),
                                         download=True)
test_dataset = torchvision.datasets.MNIST(root='../../data',
                                         train=False,
                                         transform=transforms.ToTensor())

Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Failed to download (trying next):
HTTP Error 503: Service Unavailable

Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to ../../data\MNIST\raw\train-images-idx3-ubyte.gz


31.1%IOPub message rate exceeded.
The notebook server will temporarily stop sending output
to the client in order to avoid crashing it.
To change this limit, set the config variable
`--NotebookApp.iopub_msg_rate_limit`.

Current values:
NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)
NotebookApp.rate_limit_window=3.0 (secs)

99.0%IOPub message rate exceeded.
The notebook server will temporarily stop sending output
to the client in order to avoid crashing it.
To change this limit, set the config variable
`--NotebookApp.iopub_msg_rate_limit`.

Current values:
NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)
NotebookApp.rate_limit_window=3.0 (secs)

102.8%


Extracting ../../data\MNIST\raw\train-labels-idx1-ubyte.gz to ../../data\MNIST\raw

Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ../../data\MNIST\raw\t10k-images-idx3-ubyte.gz


100.0%


Extracting ../../data\MNIST\raw\t10k-images-idx3-ubyte.gz to ../../data\MNIST\raw

Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ../../data\MNIST\raw\t10k-labels-idx1-ubyte.gz


112.7%

Extracting ../../data\MNIST\raw\t10k-labels-idx1-ubyte.gz to ../../data\MNIST\raw




d:\programdata\anaconda3\envs\rl36\lib\site-packages\torchvision\datasets\mnist.py:498: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  ..\torch\csrc\utils\tensor_numpy.cpp:180.)
  return torch.from_numpy(parsed.astype(m[2], copy=False)).view(*s)

# Data Loader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                          batch_size=batch_size,
                                          shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                         batch_size=batch_size,
                                         shuffle=False)

# 全连接神经网络含有一个隐藏层
class NeuralNet(nn.Module):
    def __init__(self,input_size,hidden_size,num_classes):
        super(NeuralNet,self).__init__()
        self.fc1 = nn.Linear(input_size,hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size,num_classes)
    
    def forward(self,x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

model = NeuralNet(input_size,hidden_size,num_classes).to(device)

# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=learning_rate)

# 训练模型
total_step = len(train_loader)
for epoch in range(num_epochs):
    for i,(images,labels) in enumerate(train_loader):
        # 将tensor添至已配置的设备
        images = images.reshape(-1,28*28).to(device)
        labels = labels.to(device)
        
        # 前向传播
        outputs = model(images)
        loss = criterion(outputs,labels)
        
        # 后向传播并优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

Epoch [1/5], Step [100/600], Loss: 0.3733
Epoch [1/5], Step [200/600], Loss: 0.2066
Epoch [1/5], Step [300/600], Loss: 0.2734
Epoch [1/5], Step [400/600], Loss: 0.1230
Epoch [1/5], Step [500/600], Loss: 0.1485
Epoch [1/5], Step [600/600], Loss: 0.0852
Epoch [2/5], Step [100/600], Loss: 0.1398
Epoch [2/5], Step [200/600], Loss: 0.1177
Epoch [2/5], Step [300/600], Loss: 0.0980
Epoch [2/5], Step [400/600], Loss: 0.0938
Epoch [2/5], Step [500/600], Loss: 0.0921
Epoch [2/5], Step [600/600], Loss: 0.0716
Epoch [3/5], Step [100/600], Loss: 0.0264
Epoch [3/5], Step [200/600], Loss: 0.1011
Epoch [3/5], Step [300/600], Loss: 0.0562
Epoch [3/5], Step [400/600], Loss: 0.0204
Epoch [3/5], Step [500/600], Loss: 0.0403
Epoch [3/5], Step [600/600], Loss: 0.1150
Epoch [4/5], Step [100/600], Loss: 0.0869
Epoch [4/5], Step [200/600], Loss: 0.0549
Epoch [4/5], Step [300/600], Loss: 0.0449
Epoch [4/5], Step [400/600], Loss: 0.0930
Epoch [4/5], Step [500/600], Loss: 0.0387
Epoch [4/5], Step [600/600], Loss: 0.0083
Epoch [5/5], Step [100/600], Loss: 0.0543
Epoch [5/5], Step [200/600], Loss: 0.0432
Epoch [5/5], Step [300/600], Loss: 0.0420
Epoch [5/5], Step [400/600], Loss: 0.0528
Epoch [5/5], Step [500/600], Loss: 0.0448
Epoch [5/5], Step [600/600], Loss: 0.0292

# 测试模型
# 测试阶段，不需要计算梯度下降
with torch.no_grad():
    correct = 0
    total = 0
    for images,labels in test_loader:
        images = images.reshape(-1,28*28).to(device)
        labels = labels.to(device)
        outputs = 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: {} %'.format(100 * correct / total))

Accuracy of the network on the 10000 test images: 97.55 %

# 保存模型为checkpoint格式
torch.save(model.state_dict(),'model.ckpt')