数据集官网:http://yann.lecun.com/exdb/mnist/
代码包含自动下载数据集
import torch
from torch import nn,optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets
import os
batch_size = 200 #分批训练数据、每批数据量
learning_rate = 1e-2 #学习率
num_epoches = 50 #训练次数
DOWNLOAD_MNIST = False #是否网上下载数据
# Mnist digits dataset
if not(os.path.exists('./mnist/')) or not os.listdir('./mnist/'):
# not mnist dir or mnist is empyt dir
DOWNLOAD_MNIST = True
train_dataset = datasets.MNIST(
root = './mnist',
train= True, #download train data
transform = transforms.ToTensor(),
download=DOWNLOAD_MNIST
)
test_dataset = datasets.MNIST(
root = './mnist',
train= False, #download test data
transform = transforms.ToTensor(),
download=DOWNLOAD_MNIST
)
#该接口主要用来将自定义的数据读取接口的输出或者PyTorch已有的数据读取接口的输入
# 按照batch size封装成Tensor,后续只需要再包装成Variable即可作为模型的输入
train_loader = DataLoader(train_dataset, batch_size = batch_size,shuffle=True) #shuffle 是否打乱加载数据
test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=False)
class Neuralnetwork(nn.Module):
def __init__(self,in_dim,n_hidden_1,n_hidden_2,out_dim):
super(Neuralnetwork,self).__init__()
self.layer1 = nn.Linear(in_dim,n_hidden_1)
self.layer2 = nn.Linear(n_hidden_1,n_hidden_2)
self.layer3 = nn.Linear(n_hidden_2,out_dim)
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
return x
model = Neuralnetwork(28*28, 300,100,10) #将28*28像素的图像二维数据展开成一维向量作为神经网络的输入
if torch.cuda.is_available(): #是否可用GPU计算
model = model.cuda() #转换成可用GPU计算的模型
criterion = nn.CrossEntropyLoss() #多分类用的交叉熵损失函数
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
#常用优化方法有
#1.Stochastic Gradient Descent (SGD)
#2.Momentum
#3.AdaGrad
#4.RMSProp
#5.Adam (momentum+adaGrad) 效果较好
for epoch in range(num_epoches):
print('epoch {}'.format(epoch+1))
print('*' * 10)
running_loss = 0.0
running_acc = 0.0 #准确度
#训练
for i,data in enumerate(train_loader,1):
img, label = data
img = img.view(img.size(0),-1) #view()修改维度,即修改shape
#判断是否可以使用GPU,若可以则将数据转化为GPU可以处理的格式。
if torch.cuda.is_available():
img = Variable(img).cuda()
label = Variable(label).cuda()
else:
img = Variable(img)
label = Variable(label)
out = model(img)
loss = criterion(out,label) #计算预测结果out和实际结果label的误差损失,(out为每个预测分类的概率)
running_loss += loss.item() * label.size(0) #先计算总的损失
_,pred = torch.max(out,1) #torch.max返回指定维度(1)中的最大值和相应序号,所以pred为预测的分类
num_correct = (pred==label).sum()
running_acc +=num_correct.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('Finish {} epoch,Loss:{:.6f},Acc:{:.6f}'.format(
epoch+1,running_loss/(len(train_dataset)),running_acc/len(train_dataset)
))
#测试
model.eval() #eval()时,模型会自动把BN和DropOut固定住,不会取平均,而是用训练好的值
eval_loss = 0.
eval_acc = 0.
for data in test_loader:
img, label = data
img = img.view(img.size(0),-1)
if torch.cuda.is_available():
img = Variable(img).cuda()
label = Variable(label).cuda()
else:
img = Variable(img)
label = Variable(label)
out = model(img)
loss = criterion(out,label)
eval_loss += loss.item() * label.size(0)
_, pred = torch.max(out, 1)
num_correct = (pred == label).sum()
eval_acc += num_correct.item()
print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len(
test_dataset)), eval_acc/len(test_dataset)))