[PyTorch小试牛刀]实战三·DNN实现逻辑回归对FashionMNIST数据集进行分类
内容还包括了网络模型参数的保存于加载。
数据集
下载地址
代码部分
import torch as t
import torchvision as tv
import numpy as np
# 超参数
EPOCH = 10
BATCH_SIZE = 100
DOWNLOAD_MNIST = True # 下过数据的话, 就可以设置成 False
N_TEST_IMG = 10 # 到时候显示 5张图片看效果, 如上图一
class DNN(t.nn.Module):
def __init__(self):
super(DNN, self).__init__()
train_data = tv.datasets.FashionMNIST(
root="./mnist/",
train=True,
transform=tv.transforms.ToTensor(),
download=DOWNLOAD_MNIST
)
test_data = tv.datasets.FashionMNIST(
root="./mnist/",
train=False,
transform=tv.transforms.ToTensor(),
download=DOWNLOAD_MNIST
)
print(test_data)
# Data Loader for easy mini-batch return in training, the image batch shape will be (50, 1, 28, 28)
self.train_loader = t.utils.data.DataLoader(
dataset=train_data,
batch_size=BATCH_SIZE,
shuffle=True)
self.test_loader = t.utils.data.DataLoader(
dataset=test_data,
batch_size=1000,
shuffle=True)
self.dnn = t.nn.Sequential(
t.nn.Linear(28*28,256),
t.nn.Dropout(0.5),
t.nn.ELU(),
t.nn.Linear(256,64),
t.nn.Dropout(0.5),
t.nn.ELU(),
t.nn.Linear(64,10)
)
self.lr = 0.001
self.loss = t.nn.CrossEntropyLoss()
self.opt = t.optim.Adam(self.parameters(), lr = self.lr)
def forward(self,x):
out = self.dnn(x)
return(out)
def train():
model = DNN()
print(model)
loss = model.loss
opt = model.opt
dataloader = model.train_loader
testloader = model.test_loader
for e in range(EPOCH):
step = 0
for (x, y) in (dataloader):
model.train()# train model dropout used
step += 1
b_x = x.view(-1, 28*28) # batch x, shape (batch, 28*28)
b_y = y
out = model(b_x)
losses = loss(out,b_y)
opt.zero_grad()
losses.backward()
opt.step()
if(step%100 == 0):
print(e,step,losses.data.numpy())
model.eval() # train model dropout not use
for (tx,ty) in testloader:
t_x = tx.view(-1, 28*28) # batch x, shape (batch, 28*28)
t_y = ty
t_out = model(t_x)
acc = (np.argmax(t_out.data.numpy(),axis=1) == t_y.data.numpy())
print(np.sum(acc)/1000)
break#只测试前1000个
t.save(model, './model.pkl') # 保存整个网络
t.save(model.state_dict(), './model_params.pkl') # 只保存网络中的参数 (速度快, 占内存少)
#加载参数的方式
"""net = DNN()
net.load_state_dict(t.load('./model_params.pkl'))
net.eval()"""
#加载整个模型的方式
net = t.load('./model.pkl')
net.eval()
for (tx,ty) in testloader:
t_x = tx.view(-1, 28*28) # batch x, shape (batch, 28*28)
t_y = ty
t_out = net(t_x)
acc = (np.argmax(t_out.data.numpy(),axis=1) == t_y.data.numpy())
print(np.sum(acc)/1000)
if __name__ == "__main__":
train()
输出结果
9 500 0.42454192
0.875
9 600 0.4553349
0.888
0.876
0.868
0.868
0.881
0.864
0.87
0.87
0.854
0.871
0.879