通过PyTorch构建的LeNet-5网络对手写数字进行训练和识别

转x33g5p2x  于2022-06-27 转载在 其他  
字(8.7k)|赞(0)|评价(0)|浏览(470)

    调用PyTorch相关接口实现一个LeNet-5网络，然后通过MNIST数据集训练模型，最后对生成的模型进行预测，主要包括2大部分：训练和预测

    1.训练部分：

    (1).加载MNIST数据集，通过调用TorchVision模块中的接口实现，将每幅图像缩放到32*32大小，小批量数据集数量设置为32；

    (2).设置网络参数的初始值，这样保证每次重新训练时初始值都是固定的，便于查找定位问题；

    (3).设计LeNet-5网络，并实例化一个网络对象，重载了__init__和forward两个函数，使用到的layer包括Conv2d、AvgPool2d、Linear；激活函数使用Tanh：

    (4).指定优化算法，这里采用Adam；

    (5).指定损失函数，这里采用CrossEntropyLoss；

    (6).训练，epochs设置为10，给出每次的训练结果；

    (7).保存模型，推荐使用state_dict。

    代码段如下：

  1. def load_mnist_dataset(img_size, batch_size):
  2.     '''下载并加载mnist数据集
  3.         img_size: 图像大小,宽高长度相同
  4.         batch_size: 小批量数据集数量
  5.     '''
  7.     # 对PIL图像先进行缩放操作,然后转换成tensor类型
  8.     transforms_ = transforms.Compose([transforms.Resize(size=(img_size, img_size)), transforms.ToTensor()])
  10.     '''下载MNIST数据集
  11.         root: mnist数据集存放目录名
  12.         train: 可选参数, 默认为True; 若为True,则从MNIST/processed/training.pt创建数据集;若为False,则从MNIST/processed/test.pt创建数据集
  13.         transform: 可选参数, 默认为None; 接收PIL图像并作处理
  14.         target_transform: 可选参数, 默认为None
  15.         download: 可选参数, 默认为False; 若为True,则从网络上下载数据集到root指定的目录
  16.     '''
  17.     train_dataset = datasets.MNIST(root="mnist_data", train=True, transform=transforms_, target_transform=None, download=True)
  18.     valid_dataset = datasets.MNIST(root="mnist_data", train=False, transform=transforms_, target_transform=None, download=False)
  20.     # 加载MNIST数据集:shuffle为True,则在每次epoch时重新打乱顺序
  21.     train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
  22.     valid_loader = DataLoader(dataset=valid_dataset, batch_size=batch_size, shuffle=False)
  24.     return train_loader, valid_loader, train_dataset, valid_dataset
  26. class LeNet5(nn.Module):
  27.     '''构建lenet网络'''
  29.     def __init__(self, n_classes: int) -> None:
  30.         super(LeNet5, self).__init__() # 调用父类Module的构造方法
  31.         # n_classes: 类别数
  33.         # nn.Sequential: 顺序容器,Module将按照它们在构造函数中传递的顺序添加,它允许将整个容器视为单个module
  34.         self.feature_extractor = nn.Sequential( # 输入32*32
  35.             nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=0), # 卷积层,28*28*6
  36.             nn.Tanh(), # 激活函数Tanh,使其值范围在(-1, 1)内
  37.             nn.AvgPool2d(kernel_size=2, stride=None, padding=0), # 平均池化层,14*14*6
  38.             nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0), # 10*10*16
  39.             nn.Tanh(),
  40.             nn.AvgPool2d(kernel_size=2, stride=None, padding=0), # 5*5*16
  41.             nn.Conv2d(in_channels=16, out_channels=120, kernel_size=5, stride=1, padding=0), # 1*1*120
  42.             nn.Tanh()
  43.         )
  45.         self.classifier = nn.Sequential( # 输入1*1*120
  46.             nn.Linear(in_features=120, out_features=84), # 全连接层,84
  47.             nn.Tanh(),
  48.             nn.Linear(in_features=84, out_features=n_classes) # 10
  49.         )
  51.     # LeNet5继承nn.Module,定义forward函数后,backward函数就会利用Autograd被自动实现
  52.     # 只要实例化一个LeNet5对象并传入对应的参数x就可以自动调用forward函数
  53.     def forward(self, x: Tensor):                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              
  54.         x = self.feature_extractor(x)
  55.         x = torch.flatten(input=x, start_dim=1) # 将输入按指定展平,start_dim=1则第一维度不变,后面的展平
  56.         logits = self.classifier(x)
  57.         probs = F.softmax(input=logits, dim=1) # 激活函数softmax: 使得每一个元素的范围都在(0,1)之间，并且所有元素的和为1
  58.         return logits, probs
  60. def validate(valid_loader, model, criterion, device):
  61.     '''Function for the validation step of the training loop'''
  63.     model.eval() # 将网络设置为评估模式
  64.     running_loss = 0
  66.     for X, y_true in valid_loader:
  67.         X = X.to(device) # 将数据导入到指定的设备上(cpu或gpu)
  68.         y_true = y_true.to(device)
  70.         # Forward pass and record loss
  71.         y_hat, _ = model(X) # 前向传播:调用Module的__call__方法, 此方法内会调用指定网络(如LeNet5)的forward方法
  72.         loss = criterion(y_hat, y_true) # 计算loss,同上,通过__call__方法调用指定损失函数类(如CrossEntropyLoss)中的forward方法
  73.         running_loss += loss.item() * X.size(0)
  75.     epoch_loss = running_loss / len(valid_loader.dataset)
  76.     return model, epoch_loss
  78. def get_accuracy(model, data_loader, device):
  79.     '''Function for computing the accuracy of the predictions over the entire data_loader'''
  81.     correct_pred = 0
  82.     n = 0
  84.     with torch.no_grad(): # 临时将循环内的所有Tensor的requires_grad标志设置为False,不再计算Tensor的梯度(自动求导)
  85.         model.eval() # 将网络设置为评估模式
  86.         for X, y_true in data_loader:
  88.             X = X.to(device) # 将数据导入到指定的设备上(cpu或gpu)
  89.             y_true = y_true.to(device)
  91.             _, y_prob = model(X) # y_prob.size(): troch.Size([32, 10]): [cols, rows]
  92.              # torch.max(input):返回Tensor中所有元素的最大值
  93.              # torch.max(input, dim):按维度dim返回最大值,并且返回索引
  94.              # dim=0: 返回每一列中最大值的那个元素,并且返回索引
  95.              # dim=1: 返回每一行中最大值的那个元素,并且返回索引
  96.             _, predicted_labels = torch.max(y_prob, 1)
  98.             n += y_true.size(0)
  99.             correct_pred += (predicted_labels == y_true).sum()
  101.     return correct_pred.float() / n
  103. def train(train_loader, model, criterion, optimizer, device):
  104.     '''Function for the training step of the training loop'''
  106.     model.train() # 将网络设置为训练模式
  107.     running_loss = 0
  109.     for X, y_true in train_loader: # 先调用DataLoader类的__iter__函数,接着循环调用_DataLoaderIter类的__next__函数
  110.         # X.size(shape: [n,c,h,w]): torch.Size([32, 1, 32, 32]); y_true.size: torch.Size([32]); n为batch_size
  111.         optimizer.zero_grad() # 将优化算法中的梯度重置为0,需要在计算下一个小批量数据集的梯度之前调用它,否则梯度将累积到现有的梯度中
  113.         # 将Tensor数据导入到指定的设备上(cpu或gpu)
  114.         X = X.to(device)
  115.         y_true = y_true.to(device)
  117.         y_hat, _ = model(X) # 前向传播:调用Module的__call__方法, 此方法内会调用指定网络(如LeNet5)的forward方法
  118.         # y_hat.size(): torch.Size([32, 10]); _.size(): torch.Size([32, 10])
  119.         loss = criterion(y_hat, y_true) # 计算loss,同上,通过__call__方法调用指定损失函数类(如CrossEntropyLoss)中的forward方法
  120.         running_loss += loss.item() * X.size(0)
  122.         loss.backward() # 反向传播,使用Autograd自动计算标量的当前梯度
  123.         optimizer.step() # 根据梯度更新网络参数,优化器通过.grad中存储的梯度来调整每个参数
  125.     epoch_loss = running_loss / len(train_loader.dataset)
  126.     return model, optimizer, epoch_loss
  128. def training_loop(model, criterion, optimizer, train_loader, valid_loader, epochs, device, print_every=1):
  129.     '''Function defining the entire training loop
  130.         model: 网络对象
  131.         criterion: 损失函数对象
  132.         optimizer: 优化算法对象
  133.         train_loader: 训练数据集对象
  134.         valid_loader: 测试数据集对象
  135.         epochs: 重复训练整个训练数据集的次数
  136.         device: 指定在cpu上还是在gpu上运行
  137.         print_every: 每训练几次打印一次训练结果
  138.     '''
  140.     train_losses = []
  141.     valid_losses = []
  143.     for epoch in range(0, epochs):
  144.         model, optimizer, train_loss = train(train_loader, model, criterion, optimizer, device)
  145.         train_losses.append(train_loss)
  147.         # 每次训练完后通过测试数据集进行评估
  148.         with torch.no_grad(): # 临时将循环内的所有Tensor的requires_grad标志设置为False,不再计算Tensor的梯度(自动求导)
  149.             model, valid_loss = validate(valid_loader, model, criterion, device)
  150.             valid_losses.append(valid_loss)
  152.         if epoch % print_every == (print_every - 1):
  153.             train_acc = get_accuracy(model, train_loader, device=device)
  154.             valid_acc = get_accuracy(model, valid_loader, device=device)
  156.             print(f'  {datetime.now().time().replace(microsecond=0)}:'
  157.                   f' Epoch: {epoch}', f' Train loss: {train_loss:.4f}', f' Valid loss: {valid_loss:.4f}'
  158.                   f' Train accuracy: {100 * train_acc:.2f}', f' Valid accuracy: {100 * valid_acc:.2f}')
  160.     return model, optimizer, (train_losses, valid_losses)
  162. def train_and_save_model():
  163.     print("#### start training ... ####")
  164.     print("1. load mnist dataset")
  165.     train_loader, valid_loader, _, _ = load_mnist_dataset(img_size=32, batch_size=32)
  167.     print("2. fixed random init value")
  168.     # 用于设置随机初始化;如果不设置每次训练时的网络初始化都是随机的,导致结果不确定;如果设置了,则每次初始化都是固定的
  169.     torch.manual_seed(seed=42)
  170.     #print("value:", torch.rand(1), torch.rand(1), torch.rand(1)) # 运行多次,每次输出的值都是相同的,[0, 1)
  172.     print("3. instantiate lenet net object")
  173.     model = LeNet5(n_classes=10).to('cpu') # 在CPU上运行
  174.     print("4. specify the optimization algorithm: Adam")
  175.     optimizer = torch.optim.Adam(params=model.parameters(), lr=0.001) # 定义优化算法:Adam是一种基于梯度的优化算法
  176.     print("5. specify the loss function: CrossEntropyLoss")
  177.     criterion = nn.CrossEntropyLoss() # 定义损失函数:交叉熵损失
  179.     print("6. repeated training")
  180.     model, _, _ = training_loop(model, criterion, optimizer, train_loader, valid_loader, epochs=10, device='cpu') # epochs为遍历训练整个数据集的次数
  182.     print("7. save model")
  183.     model_name = "../../../data/Lenet-5.pth"
  184.     #torch.save(model, model_name) # 保存整个模型, 对应于model = torch.load
  185.     torch.save(model.state_dict(), model_name) # 推荐:只保存模型训练好的参数,对应于model.load_state_dict(torch.load)

    执行结果如下所示：

    2.手写数字图像识别部分：

    (1).加载模型，推荐使用load_state_dict，对应于保存模型时使用的state_dict；

    (2).设置网络到评估模式；

    (3).准备测试图像，一共10幅，0到9各一幅，如下图所示，注意：训练图像背景色为黑色，而测试图像背景色为白色：

    (4).依次对每幅图像进行识别。

    代码段如下所示：

  1. def list_files(filepath, filetype):
  2.     '''遍历指定目录下的指定文件'''
  4.     paths = []
  5.     for root, dirs, files in os.walk(filepath):
  6.         for file in files:
  7.             if file.lower().endswith(filetype.lower()):
  8.                 paths.append(os.path.join(root, file))
  9.     return paths
  11. def get_image_label(image_name, image_name_suffix):
  12.     '''获取测试图像对应label'''
  14.     index = image_name.rfind("/")
  15.     if index == -1:
  16.         print(f"Error: image name {image_name} is not supported")
  18.     sub = image_name[index+1:]
  19.     label = sub[:len(sub)-len(image_name_suffix)]
  20.     return label
  22. def image_predict():
  23.     print("#### start predicting ... ####")
  24.     print("1. load model")
  25.     model_name = "../../../data/Lenet-5.pth"
  26.     model = LeNet5(n_classes=10).to('cpu') # 实例化一个网络对象
  27.     model.load_state_dict(torch.load(model_name)) # 加载模型
  29.     print("2. set net to evaluate mode")
  30.     model.eval()
  32.     print("3. prepare test images")
  33.     image_path = "../../../data/image/handwritten_digits/"
  34.     image_name_suffix = ".png"
  35.     images_name = list_files(image_path, image_name_suffix)
  37.     print("4. image recognition")
  38.     with torch.no_grad():
  39.         for image_name in images_name:
  40.             #print("image name:", image_name)
  41.             label = get_image_label(image_name, image_name_suffix)
  43.             img = cv2.imread(image_name, cv2.IMREAD_GRAYSCALE)
  44.             img = cv2.resize(img, (32, 32))
  45.             # MNIST图像背景为黑色,而测试图像的背景色为白色,识别前需要做转换
  46.             img = cv2.bitwise_not(img)
  47.             #print("img shape:", img.shape)
  49.             # 将opencv image转换到pytorch tensor
  50.             transform = transforms.ToTensor()
  51.             tensor = transform(img) # tensor shape: torch.Size([1, 32, 32])
  52.             tensor = tensor.unsqueeze(0) # tensor shape: torch.Size([1, 1, 32, 32])
  53.             #print("tensor shape:", tensor.shape)
  55.             _, y_prob = model(tensor)
  56.             _, predicted_label = torch.max(y_prob, 1)
  57.             print(f"  predicted label: {predicted_label.item()}, ground truth label: {label}")

    执行结果如下图所示：

    GitHub： https://github.com/fengbingchun/PyTorch_Test

版权说明 : 本文为转载文章, 版权归原作者所有 版权申明

原文链接 : https://blog.csdn.net/fengbingchun/article/details/125462001

内容来源于网络，如有侵权，请联系作者删除！

pytorch

相关文章

    查看更多

    热门标签

    更多
    JavaquerypythonNode开发语言requestUtil数据库Table后端算法LoggerMessageElementParser

    最新文章

    更多
    • 蜀ICP备13028337号-1 大数据知识库 https://www.saoniuhuo.com © All rights reserved
    • 本站内容来源互联网,如果侵犯您的权益请联系我们删除, 联系方式：448109455@qq.com