MNIST手写识别

import numpy as np
# import matplotlib.pyplot as plt
import torch
from torch import nn,optim
from torch.autograd import Variable
# import torch
from torchvision import datasets,transforms
from torch.utils.data import DataLoader

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1=nn.Linear(784,10)
        self.softmax=nn.Softmax(dim=1)

    def forward(self,x):
        x=x.view(x.size()[0],-1)
        x=self.fc1(x)
        x=self.softmax(x)
        return x


#载入训练集
train_dataset=datasets.MNIST(root='./',train=True,transform=transforms.ToTensor(),download=True)

test_dataset=datasets.MNIST(root='./',train=False,transform=transforms.ToTensor(),download=True)

batch_size=64

train_loader=DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True)


test_loader=DataLoader(dataset=test_dataset,batch_size=batch_size,shuffle=True)



LR=0.5
model=Net()
mse_loss=nn.MSELoss()
optimizer=optim.SGD(model.parameters(),LR)


def train():
    for i, data in enumerate(train_loader):
        inputs, labels = data#输入数据和真正的标签
        out=model(inputs)
        labels=labels.reshape(-1,1)
        one_hot=torch.zeros(inputs.shape[0],10).scatter(1,labels,1)
        #计算loss
        loss=mse_loss(out,one_hot)
        #梯度清0
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


def test():
    correct=0
    for i,data in enumerate(test_loader):
        inputs,labels=data

        out=model(inputs)
        _,predicted=torch.max(out,1)
        correct+=(predicted==labels).sum()
    print('test acc:{0}'.format(correct/len(test_dataset)))

for epoch in range(100):
    print('epoch',epoch)
    train()
    test()

效果：