运行代码1为:
import numpy as np
from torch.nn import Module
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from matplotlib import pyplot as plt
import torch
import torch.nn as nn
import torchvision
from tqdm import tqdm
mnist = MNIST(root='~', train=True, download=True)
trans = torchvision.transforms.Compose(
[torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.1307,), (0.3081,))])
trainData = torchvision.datasets.MNIST(root='~', train=True, transform=trans, download=True)
testData = torchvision.datasets.MNIST(root='~', train=False, transform=trans, download=True)
trainDataLoader = DataLoader(dataset=trainData, batch_size=16, shuffle=True)
testDataLoader = DataLoader(dataset=testData, batch_size=16, shuffle=True)
class Model(nn.Module):
def __init__(self):
super().__init__()
self.inputlayer = nn.Sequential(nn.Linear(28 * 28, 256), nn.ReLU(), nn.Dropout(0.2))
self.hiddenlayer = nn.Sequential(nn.Linear(256, 256), nn.ReLU(), nn.Dropout(0.2))
self.outlayer = nn.Sequential(nn.Linear(256, 10))
def forward(self, x):
x = x.view(x.size(0), -1)
x = self.inputlayer(x)
x = self.hiddenlayer(x)
x = self.outlayer(x)
return x
model = Model()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
celoss = nn.CrossEntropyLoss()
best_acc = 0
correct = 0
loss_all = {'train': [], 'Test ': []}
for epoch in tqdm(range(3)):
processBar = tqdm(trainDataLoader, unit='step')
model.eval()
numer_val, denumer_val, loss_tr = 0., 0., 0.
with torch.no_grad():
for data, targe in testDataLoader:
output = model(data)
loss = celoss(output, targe)
loss_tr += loss.data
model.train()
loss_train = 0.
for imag, labels in trainDataLoader:
optimizer.zero_grad()
output = model(data)
loss = celoss(output, labels)
loss_train += loss.data
loss.backward()
optimizer.step()
loss_all['train'].append(loss_tr / len(trainDataLoader))
loss_all['Test '].append(loss_train / len(testDataLoader))
plt.plot(loss_all['train'])
plt.plot(loss_all['Test '])
plt.show()
运行结果:
网上有文章讲述的特别好,代码也比这个高级很多,详情:https://blog.csdn.net/NikkiElwin/article/details/112980305
运行代码2
import torch
import torch.nn as nn
import torch.nn.functional as functional
import torchvision
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from PIL import Image
from copy import deepcopy
import matplotlib.pyplot as plt
import numpy as np
from torch.utils.data import DataLoader
mnist_data = MNIST(root='~', train=True, download=True) # 图片的shape是1*28*28
# print(mnist_data[0][0])
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 20, 5, 1)
# torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1)
# in_channels:输入图像通道数,手写数字图像为1,彩色图像为3
# out_channels:输出通道数,这个等于卷积核的数量
# kernel_size:卷积核大小
# stride:步长
self.conv2 = nn.Conv2d(20, 50, 5, 1)
self.fc1 = nn.Linear(4 * 4 * 50, 500)
self.fc2 = nn.Linear(500, 10)
# 最后一个全连接层的输出维度必须是分类的类数,所以是10
def forward(self, x):
x = functional.relu(self.conv1(x)) # 变成20*24*24(24=28+1-5)
x = functional.max_pool2d(x, 2, 2) # 变成20*12*12
x = functional.relu(self.conv2(x)) # 变成50*8*8
x = functional.max_pool2d(x, 2, 2) # 变成50*4*4
x = x.view(-1, 4 * 4 * 50)
#注意一点,在输入到全连接层之前,需要将其从一个三维的张量转换成一个一维张量,用到view()
x = functional.relu(self.fc1(x))
x = self.fc2(x)
return functional.log_softmax(x, dim=1)
device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
batch_size = 32
train_dataloder = DataLoader(
MNIST(root='~', train=True, download=True, transform=trans) ,batch_size=batch_size, shuffle=True)
test_dataloder = DataLoader(
MNIST(root='~', train=False, download=True, transform=trans), batch_size=batch_size, shuffle=True)
lr = 0.01
momentum = 0.5
model = Net().to(device) # 模型初始化
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum) # 定义优化器
for epoch in range(10):
model.train()
for batch_idx, (data, target) in enumerate(train_dataloder):
data, target = data.to(device), target.to(device)
pred = model(data)
loss = functional.nll_loss(pred, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % 100 == 0:
print("Train Epoch: {} [{}/{} ({:0f}%)]\tLoss: {:.6f}".format(
epoch,
batch_idx * len(data), # 100*32
len(train_dataloder.dataset), # 60000
100. * batch_idx / len(train_dataloder), # len(train_loader)=60000/32=1875
loss.item()
))
model.eval()
total_loss=0.
correct=0.
with torch.no_grad():
for data, target in test_dataloder:
data, target = data.to(device), target.to(device)
output = model(data)
pred = model(data).argmax(dim=1)#找出最高的那个,就是预测的结果
total_loss+= functional.nll_loss(output, target,reduction='sum').item()
correct+=pred.eq(target.view_as(pred)).sum().item()
total_loss /= len(test_dataloder.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
total_loss, correct, len(test_dataloder.dataset),
100. * correct / len(test_dataloder.dataset)))