import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(0.1307, 0.3081)])
train_Data = datasets.MNIST('D:/PyCharm/minist/', train=True
, download=True, transform=transform)
test_Data = datasets.MNIST('D:/PyCharm/minist/', train=False
, download=True, transform=transform)
train_loader = DataLoader(train_Data, batch_size=128, shuffle=True)
test_loader = DataLoader(test_Data, batch_size=64, shuffle=False)
class DNN(nn.Module):
def __init__(self):
super(DNN, self).__init__()
self.net = nn.Sequential(
nn.Linear(784, 512),
nn.ReLU(),
nn.Linear(512, 256),
nn.ReLU(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, 10)
)
def forward(self, x):
y = self.net(x)
return y
model = DNN().to('cuda:0')
print(model)
loss_fn=nn.CrossEntropyLoss()
learning_rate=0.01
optimizer=torch.optim.SGD(model.parameters(),lr=learning_rate,momentum=0.9)
epochs=5
losses=[]
for epoch in range(epochs):
for(x,y) in train_loader:
x=x.view(x.size(0),-1)
x,y=x.to('cuda:0'),y.to('cuda:0')
Pred=model(x)
loss=loss_fn(Pred,y)
losses.append(loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
Fig=plt.figure()
plt.plot(range(len(losses)),losses)
plt.show()
correct=0
total=0
for(x,y) in test_loader:
x=x.view(x.size(0),-1)
x,y=x.to('cuda:0'),y.to('cuda:0')
Pred=model(x)
_,predicted=torch.max(Pred.data,dim=1)
correct+=torch.sum((predicted==y))
total+=y.size(0)
print(f"测试集进准度:{100*correct/total}%")