import torch
import torchvision.datasets as dataset
import torch.utils.data as data_utils
from torchvision import transforms
# 训练集下载
train_data = dataset.MNIST(root='mnist', train=True, transform=transforms.ToTensor(), download=True)
# 测试集下载
test_data = dataset.MNIST(root='mnist', train=False, transform=transforms.ToTensor(), download=False)
# batchsize dataloader可以将数据进行循环读取
train_loader = data_utils.DataLoader(dataset=train_data, batch_size=64, shuffle=True)
test_loader = data_utils.DataLoader(dataset=test_data, batch_size=64, shuffle=True)
class CNN(torch.nn.Module):
def __init__(self):
super(CNN, self).__init__()
# 定义卷积层
self.conv = torch.nn.Sequential(torch.nn.Conv2d(1, 32, kernel_size=5, padding=2), torch.nn.BatchNorm2d(32),
torch.nn.ReLU(),
torch.nn.MaxPool2d(2))
# 定义网络层
self.fc = torch.nn.Linear(14 * 14 * 32, 10)
def forward(self, x):
out = self.conv(x)
out = out.view(out.size()[0], -1)
out = self.fc(out)
return out
cnn = CNN()
# loss 损失函数,采用交叉熵
loss_func = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(cnn.parameters(), lr=0.01)
for epoch in range(10):
for i, (images, labels) in enumerate(train_loader):
outputs = cnn(images)
loss = loss_func(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f'epoch is {epoch + 1}, ite is {len(train_data)}/{64}, loss is {loss.item()}')
# 保存训练好的模型文件
torch.save(cnn, 'mnist.pkl')
PyTorch小案例——手写字识别
于 2022-02-18 00:04:12 首次发布