[Pyrorch] MNIST 使用不同优化器

import torch
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.nn.functional as F    #选择激活函数
import torch.optim as optim  #选择优化器
import matplotlib.pyplot as plt

#准备数据
batch_size=64
trans=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,),(0.3081,))])

#数据载入

train_dataset=datasets.MNIST(root='../MNNIST/mnist/',train=True,download=True,transform=trans)

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

test_dataset=datasets.MNIST(root='../MNNIST/mnist/',train=False,download=True,transform=trans)

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


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.linear1 = torch.nn.Linear(784, 512)
        self.linear2 = torch.nn.Linear(512, 64)
        self.linear3 = torch.nn.Linear(64, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = F.relu(self.linear1(x))
        x = F.relu(self.linear2(x))
        x = self.linear3(x)   #最后输出不用进入线性层
        return x

model=Net()

x=[]

criterion=torch.nn.CrossEntropyLoss() #计算损失
optimizer=optim.SGD(model.parameters(),lr=0.01,momentum=0.5)  #momentum带冲量的？？？

def train(epoch):
    running_loss=0.0

    for batch_idx,data in enumerate(train_loader,0):
        input,targe=data

        optimizer.zero_grad()

        output=model(input)
        loss=criterion(output,targe)
        loss.backward()

        optimizer.step()


        running_loss+=loss.item()  #计算输出
        x.append(loss.item())

        if batch_idx%300==299:
            print("[%d,%d]  loss: %.4f" % (epoch+1,batch_idx+1,running_loss/300))
            running_loss=0

def test():
    correct=0
    total=0
    with torch.no_grad():
        for data in test_loader:
            images,labels=data
            outputs=model(images)

            _,predicted=torch.max(outputs.data,dim=1)  #数据和数据下标

            total+=labels.size(0)  # N*1 0即为数量

            correct+=(predicted==labels).sum().item()

    print("Accuracy on test set:  %d %%" % (100*correct/total))

if __name__=='__main__':
    for epoch in range(1):
        train(epoch)
        test()
        
    plt.plot(x, label='Train_loss')
    plt.title('Train_loss')
    plt.show()