import torchvision
import torchvision.datasets as datasets
import torch
import torch.nn.functional as F
import numpy as np
import random
random.seed(1)
np.random.seed(1)
torch.manual_seed(1)
train_mnist = datasets.MNIST(root = './data', train=True, download = True)
test_mnist = datasets.MNIST(root = './data', train=False, download = True)
train_X, train_y = train_mnist.data, train_mnist.targets
test_X, test_y = test_mnist.data, test_mnist.targets
train_X = train_X.float()
test_X = test_X.float()
#数据归一化
train_X = train_X / 255.0
test_X = test_X / 255.0
dim_in = 28*28
dim_hid = 128
dim_out = 10
class TwoLayerNet(torch.nn.Module):
def __init__(self, dim_in, dim_hid, dim_out):
super(TwoLayerNet, self).__init__()
#define the model architecture
self.fc1 = torch.nn.Linear(dim_in, dim_hid, bias=True)
self.fc2 = torch.nn.Linear(dim_hid, dim_out, bias=True)
def forward(self, x):
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = F.relu(x)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
#提前定义模型
model = TwoLayerNet(dim_in, dim_hid, dim_out)
#提前定义loss和优化器
loss_fun = torch.nn.NLLLoss()
eta = 0.001
optimizer = torch.optim.Adam(model.parameters(), lr=eta)
#训练模型
for i in range(10):
#Forward pass
y_pred = model(train_X)
#loss
loss = loss_fun(y_pred, train_y)
print(i, loss.item())
optimizer.zero_grad()
#backwward pass
loss.backward()
#update model parameters
optimizer.step()
#模型评估
y_pred = model(test_X)
#get the index of the max log-probability
pred = y_pred.argmax(dim=1, keepdim=True)
correct = pred.eq(test_y.view_as(pred)).sum().item()
acc = correct / test_X.shape[0]
print('准确率为{}%'.format(acc*100))
0 2.2972536087036133
1 2.231241464614868
2 2.1664140224456787
3 2.100111246109009
4 2.0312047004699707
5 1.9600239992141724
6 1.8877311944961548
7 1.8152536153793335
8 1.743047833442688
9 1.6712080240249634
[Decode error - output not utf-8]
[Finished in 4.7s]