import torch
import torch.utils.data as Data
import torch.nn as nn
from torchvision import datasets,transforms
import torchvision
from torch.autograd import Variable
import numpy as np
import matplotlib.pyplot as plt
1.加载数据
transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5,],std=[0.5,])])
data_train=datasets.MNIST(root="D:\jupyter_data", transform=transform, train=True,
download=True
)
data_test=datasets.MNIST(root="D:\jupyter_data", transform=transform, train=False)
#按批获取数据
batch = 256
train_iter = Data.DataLoader(dataset=data_train, batch_size=batch, shuffle=True)
test_iter = Data.DataLoader(data_test,batch, shuffle=True)
print(test_iter)
<torch.utils.data.dataloader.DataLoader object at 0x0000020AE9927400>
2.定义和初始化模型
#输入层784,隐藏层20个神经元,输出10个
input_num, hidden_num, output_num=784,20,10
class FlattenLayer(nn.Module): #数据处理层,将输入变为256 * 784
def __init__(self):
super(FlattenLayer, self).__init__()
def forward(self, x): # x shape: (batch, *, *, ...)
return x.view(x.shape[0], -1)
net = nn.Sequential()
net.add_module('process', FlattenLayer())
net.add_module('linear',nn.Linear(input_num,hidden_num))
net.add_module('relu',nn.ReLU())
net.add_module('linear2',nn.Linear(hidden_num,output_num))
print(net)
#初始化模型参数
for param in net.parameters():
nn.init.normal_(param, mean=0, std=0.01)
Sequential(
(process): FlattenLayer()
(linear): Linear(in_features=784, out_features=20, bias=True)
(relu): ReLU()
(linear2): Linear(in_features=20, out_features=10, bias=True)
)
3.定义损失函数和优化方法
loss = nn.CrossEntropyLoss()
optim = torch.optim.SGD(net.parameters(), lr=0.2)
4.训练模型
epoch_size = 2
for epoch in range(0, epoch_size+1):
train_acc ,n ,train_loss = 0, 0, 0
for X,y in train_iter:
y_pre = net(X)
l = loss(y_pre, y).sum()
optim.zero_grad()
l.backward()
optim.step()
train_acc += (y_pre.argmax(dim=1) == y).sum().item()
n += y.shape[0]
train_loss += l.item()
print('epoch:%d, loss:%.4f, train_acc:%.4f'%(epoch+1, train_loss/n, train_acc/n))
epoch:1, loss:0.0052, train_acc:0.5317
epoch:2, loss:0.0030, train_acc:0.7506
epoch:3, loss:0.0020, train_acc:0.8430
def test_acc(test_iter, net):
test_acc_count = 0
n = 0
for X,y in test_iter:
y_pre = net(X)
test_acc_count += (y_pre.argmax(dim=1) == y).sum().item()
n += y.shape[0]
print('test acc %.3f'% (test_acc_count/n))
test_acc(test_iter, net)
test acc 0.770
5.可视化
import collections
for X,y in train_iter:
print(X[0][0].shape)
out = X[0][0].numpy()
feature_img = np.asarray(out * 255, dtype=np.uint8)
plt.imshow(feature_img,cmap='gray')
plt.show()
y_pre = net(X)
#获取最大值索引通过转为list
pre = y_pre[0].tolist()
print("预测:", pre.index(max(pre, key = abs)))
break
torch.Size([28, 28])
预测: 2