import torch
import torchvision.datasets
from torch import nn, optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
train_data = torchvision.datasets.MNIST("./data", train=True,
transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.MNIST("./data", train=False,
transform=torchvision.transforms.ToTensor(), download=True)
train_Dataloader = DataLoader(train_data, batch_size=64, shuffle=True)
test_Dataloader = DataLoader(test_data, batch_size=64, shuffle=False)
class Lenet_5(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
)
self.layer2 = nn.Sequential(
nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
)
self.layer3 = nn.Sequential(
nn.Linear(16*5*5, 120),
nn.ReLU(),
)
self.layer4 = nn.Sequential(
nn.Linear(120, 84),
nn.ReLU(),
)
self.layer5 = nn.Linear(84, 10)
def forward(self, input):
output = self.layer1(input)
output = self.layer2(output)
# print(output.shape)
output = output.view(output.size(0), -1)
output = self.layer3(output)
# print("layer3", output.shape)
output = self.layer4(output)
output = self.layer5(output)
return output
test_data_size = len(test_data)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
lenet_5 = Lenet_5()
writer = SummaryWriter("log_lenet_5")
lenet_5.to(device)
# loss_fuc = nn.CrossEntropyLoss().to(device) #多分类问题,选择交叉熵损失函数
loss_fn = nn.CrossEntropyLoss().to(device)
optimizer = optim.SGD(lenet_5.parameters(), lr=0.001, momentum=0.9)
epoch = 10
train_step = 0
total_test_step = 0
max_accuracy = 0
for i in range(epoch):
for data in train_Dataloader:
imgs, targets = data
if torch.cuda.is_available():
imgs = imgs.to(device)
targets = targets.to(device)
# 梯度清零
optimizer.zero_grad()
# 传递损失 + 更新参数
output = lenet_5(imgs)
# print("output_shape", output.shape)
# print("targets_shape", targets.shape)
loss = loss_fn(output, targets)
loss.backward()
optimizer.step()
if train_step % 100 ==0:
print("训练 {} 次:".format(train_step+1)+". loss={}:".format(loss))
writer.add_scalar("train_loss", loss.item(), train_step)
train_step = train_step + 1
if train_step == 3:
writer.add_graph(lenet_5, imgs)
total_accuracy = 0
total_test_loss = 0
accuracy = 0
for data in test_Dataloader:
imgs, targets = data
if torch.cuda.is_available():
imgs = imgs.to(device)
targets = targets.to(device)
outputs = lenet_5(imgs)
loss = loss_fn(outputs, targets)
total_test_loss = total_test_loss + loss.item()
accuracy = (outputs.argmax(1) == targets.reshape(targets.shape, 1)).sum().to(device)
# print(outputs.argmax(1))
# print(targets.reshape(targets.shape, 1))
total_accuracy = total_accuracy + accuracy
print("整体测试集上的正确率{}".format(total_accuracy / test_data_size))
writer.add_scalar("test_loss", total_test_loss, total_test_step)
writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)
total_test_step = total_test_step + 1
if max_accuracy < total_accuracy:
# 保存模型
torch.save(lenet_5.state_dict(), 'model.pkl')
max_accuracy = total_accuracy
writer.close()
基于Lenet-5的MNIST数据集处理
最新推荐文章于 2023-03-16 14:18:03 发布