• 基于ResNet34的花朵分类

    一.数据集准备

    新建一个项目文件夹ResNet，并在里面建立data_set文件夹用来保存数据集，在data_set文件夹下创建新文件夹"flower_data"，点击链接下载花分类数据集https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz，会下载一个压缩包，将它解压到flower_data文件夹下，执行"split_data.py"脚本自动将数据集划分成训练集train和验证集val。

     split.py如下：

    1. import os
    2. from shutil import copy, rmtree
    3. import random
    4.  
    5.  
    6. def mk_file(file_path: str):
    7.     if os.path.exists(file_path):
    8.         # 如果文件夹存在，则先删除原文件夹在重新创建
    9.         rmtree(file_path)
    10.     os.makedirs(file_path)
    11.  
    12.  
    13. def main():
    14.     # 保证随机可复现
    15.     random.seed(0)
    16.  
    17.     # 将数据集中10%的数据划分到验证集中
    18.     split_rate = 0.1
    19.  
    20.     # 指向你解压后的flower_photos文件夹
    21.     cwd = os.getcwd()
    22.     data_root = os.path.join(cwd, "flower_data")
    23.     origin_flower_path = os.path.join(data_root, "flower_photos")
    24.     assert os.path.exists(origin_flower_path), "path '{}' does not exist.".format(origin_flower_path)
    25.  
    26.     flower_class = [cla for cla in os.listdir(origin_flower_path)
    27.                     if os.path.isdir(os.path.join(origin_flower_path, cla))]
    28.  
    29.     # 建立保存训练集的文件夹
    30.     train_root = os.path.join(data_root, "train")
    31.     mk_file(train_root)
    32.     for cla in flower_class:
    33.         # 建立每个类别对应的文件夹
    34.         mk_file(os.path.join(train_root, cla))
    35.  
    36.     # 建立保存验证集的文件夹
    37.     val_root = os.path.join(data_root, "val")
    38.     mk_file(val_root)
    39.     for cla in flower_class:
    40.         # 建立每个类别对应的文件夹
    41.         mk_file(os.path.join(val_root, cla))
    42.  
    43.     for cla in flower_class:
    44.         cla_path = os.path.join(origin_flower_path, cla)
    45.         images = os.listdir(cla_path)
    46.         num = len(images)
    47.         # 随机采样验证集的索引
    48.         eval_index = random.sample(images, k=int(num*split_rate))
    49.         for index, image in enumerate(images):
    50.             if image in eval_index:
    51.                 # 将分配至验证集中的文件复制到相应目录
    52.                 image_path = os.path.join(cla_path, image)
    53.                 new_path = os.path.join(val_root, cla)
    54.                 copy(image_path, new_path)
    55.             else:
    56.                 # 将分配至训练集中的文件复制到相应目录
    57.                 image_path = os.path.join(cla_path, image)
    58.                 new_path = os.path.join(train_root, cla)
    59.                 copy(image_path, new_path)
    60.             print("\r[{}] processing [{}/{}]".format(cla, index+1, num), end="")  # processing bar
    61.         print()
    62.  
    63.     print("processing done!")
    64.  
    65.  
    66. if __name__ == '__main__':
    67.     main()

    之后会在文件夹下生成train和val数据集，到此，完成了数据集的准备。

     二.定义网络

    新建model.py，参照ResNet的网络结构和pytorch官方给出的代码，对代码进行略微的修改即可，首先定义了两个类BasicBlock和Bottleneck，分别对应着ResNet18、34和ResNet50、101、152，从下面这个图就可以区别开来。

    可见，18和34层的网络，他们的conv2_x,conv3_x,conv4_x,conv5_x是相同的，不同的是每一个block的数量（[2 2 2 2]和[3 4 6 3]），50和101和152层的网络，多了1*1卷积核，block数量也不尽相同。

    接着定义了ResNet类，进行前向传播。对于34层的网络（这里借用了知乎牧酱老哥的图，18和34的block相同，所以用18的进行讲解），conv2_x和conv3_x对应的残差块对应的残差快在右侧展示出来（可以注意一下stride），当计算特征图尺寸时，要特别注意。在下方代码计算尺寸的部分我都进行了注释。

    pytorch官方ResNet代码

    修改后的train.py：

    1. import torch.nn as nn
    2. import torch
    3.  
    4.  
    5. class BasicBlock(nn.Module):    #18 34层残差结构， 残差块
    6.     expansion = 1
    7.  
    8.     def __init__(self, in_channel, out_channel, stride=1, downsample=None, **kwargs):
    9.         super(BasicBlock, self).__init__()
    10.         self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=3, stride=stride, padding=1, bias=False)
    11.         self.bn1 = nn.BatchNorm2d(out_channel)
    12.         self.relu = nn.ReLU()
    13.         self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel, kernel_size=3, stride=1, padding=1, bias=False)
    14.         self.bn2 = nn.BatchNorm2d(out_channel)
    15.         self.downsample = downsample
    16.  
    17.     def forward(self, x):
    18.         identity = x
    19.         if self.downsample is not None:
    20.             identity = self.downsample(x)    # 不为none，对应虚线残差结构（下需要1*1卷积调整维度），为none，对应实线残差结构（不需要1*1卷积）
    21.  
    22.         out = self.conv1(x)        
    23.         out = self.bn1(out)
    24.         out = self.relu(out)
    25.  
    26.         out = self.conv2(out)     
    27.         out = self.bn2(out)
    28.  
    29.         out += identity
    30.         out = self.relu(out)
    31.  
    32.         return out
    33.  
    34.  
    35. class Bottleneck(nn.Module):    #50 101 152层残差结构
    36.     """
    37.     注意：原论文中，在虚线残差结构的主分支上，第一个1x1卷积层的步距是2，第二个3x3卷积层步距是1。
    38.     但在pytorch官方实现过程中是第一个1x1卷积层的步距是1，第二个3x3卷积层步距是2，
    39.     这么做的好处是能够在top1上提升大概0.5%的准确率。
    40.     可参考Resnet v1.5 https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch
    41.     """
    42.     expansion = 4
    43.  
    44.     def __init__(self, in_channel, out_channel, stride=1, downsample=None, groups=1, width_per_group=64):
    45.         super(Bottleneck, self).__init__()
    46.  
    47.         width = int(out_channel * (width_per_group / 64.)) * groups
    48.  
    49.         # squeeze channels
    50.         self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=width, kernel_size=1, stride=1, bias=False)  
    51.         self.bn1 = nn.BatchNorm2d(width)
    52.         # -----------------------------------------
    53.         self.conv2 = nn.Conv2d(in_channels=width, out_channels=width, groups=groups,
    54.                                kernel_size=3, stride=stride, bias=False, padding=1)
    55.         self.bn2 = nn.BatchNorm2d(width)
    56.         # -----------------------------------------
    57.         # unsqueeze channels
    58.         self.conv3 = nn.Conv2d(in_channels=width, out_channels=out_channel*self.expansion, kernel_size=1, stride=1, bias=False)  
    59.         self.bn3 = nn.BatchNorm2d(out_channel*self.expansion)
    60.         self.relu = nn.ReLU(inplace=True)
    61.         self.downsample = downsample
    62.  
    63.     def forward(self, x):
    64.         identity = x
    65.         if self.downsample is not None:  # 不为none，对应虚线残差结构（下需要1*1卷积调整维度），为none，对应实线残差结构（不需要1*1卷积）
    66.             identity = self.downsample(x)
    67.  
    68.         out = self.conv1(x)
    69.         out = self.bn1(out)
    70.         out = self.relu(out)
    71.  
    72.         out = self.conv2(out)
    73.         out = self.bn2(out)
    74.         out = self.relu(out)
    75.  
    76.         out = self.conv3(out)
    77.         out = self.bn3(out)
    78.  
    79.         out += identity
    80.         out = self.relu(out)
    81.  
    82.         return out
    83.  
    84.  
    85. class ResNet(nn.Module):
    86.  
    87.     def __init__(self, block, blocks_num, num_classes=1000, include_top=True, groups=1, width_per_group=64):
    88.         super(ResNet, self).__init__()
    89.         self.include_top = include_top
    90.         self.in_channel = 64
    91.  
    92.         self.groups = groups
    93.         self.width_per_group = width_per_group
    94.  
    95.         # （channel height width）
    96.         self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2, padding=3, bias=False) # (3 224 224) -> (64 112 112)
    97.         self.bn1 = nn.BatchNorm2d(self.in_channel)
    98.         self.relu = nn.ReLU(inplace=True)
    99.         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)         # (64 112 112) -> (64 56 56)
    100.         # 对于每一个block，第一次的两个卷积层stride=1和1，第二次stride=1和1
    101.         self.layer1 = self._make_layer(block, 64, blocks_num[0])                # (64 56 56) -> (64 56 56)
    102.         # 对于每一个block，第一次的两个卷积层stride=2和1，第二次stride=1和1
    103.         self.layer2 = self._make_layer(block, 128, blocks_num[1], stride=2)     # (64 56 56) -> (128 28 28)
    104.         self.layer3 = self._make_layer(block, 256, blocks_num[2], stride=2)     # (128 28 28) -> (256 14 14)
    105.         self.layer4 = self._make_layer(block, 512, blocks_num[3], stride=2)     # (256 28 28) -> (512 14 14)
    106.         if self.include_top:
    107.             self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # output size = (1, 1)
    108.             self.fc = nn.Linear(512 * block.expansion, num_classes)
    109.  
    110.         for m in self.modules():
    111.             if isinstance(m, nn.Conv2d):
    112.                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
    113.  
    114.     def _make_layer(self, block, channel, block_num, stride=1):           # channel为当前block所使用的卷积核个数
    115.         downsample = None 
    116.         if stride != 1 or self.in_channel != channel * block.expansion:   # 18和32不满足判断条件，会跳过；50 101 152会执行这部分
    117.             downsample = nn.Sequential(
    118.                 nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),
    119.                 nn.BatchNorm2d(channel * block.expansion))
    120.  
    121.         layers = []
    122.         layers.append(block(self.in_channel,
    123.                             channel,
    124.                             downsample=downsample,
    125.                             stride=stride,
    126.                             groups=self.groups,
    127.                             width_per_group=self.width_per_group))
    128.         self.in_channel = channel * block.expansion
    129.  
    130.         for _ in range(1, block_num):
    131.             layers.append(block(self.in_channel,
    132.                                 channel,
    133.                                 groups=self.groups,
    134.                                 width_per_group=self.width_per_group))
    135.  
    136.         return nn.Sequential(*layers)
    137.  
    138.     def forward(self, x):
    139.         x = self.conv1(x)
    140.         x = self.bn1(x)
    141.         x = self.relu(x)
    142.         x = self.maxpool(x)
    143.  
    144.         x = self.layer1(x)
    145.         x = self.layer2(x)
    146.         x = self.layer3(x)
    147.         x = self.layer4(x)
    148.  
    149.         if self.include_top:
    150.             x = self.avgpool(x)
    151.             x = torch.flatten(x, 1)
    152.             x = self.fc(x)
    153.  
    154.         return x
    155.  
    156.  
    157. def resnet34(num_classes=1000, include_top=True):
    158.     # https://download.pytorch.org/models/resnet34-333f7ec4.pth
    159.     return ResNet(BasicBlock, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top)
    160.  
    161.  
    162. def resnet50(num_classes=1000, include_top=True):
    163.     # https://download.pytorch.org/models/resnet50-19c8e357.pth
    164.     return ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top)
    165.  
    166.  
    167. def resnet101(num_classes=1000, include_top=True):
    168.     # https://download.pytorch.org/models/resnet101-5d3b4d8f.pth
    169.     return ResNet(Bottleneck, [3, 4, 23, 3], num_classes=num_classes, include_top=include_top)
    170.  
    171.  
    172. def resnext50_32x4d(num_classes=1000, include_top=True):
    173.     # https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth
    174.     groups = 32
    175.     width_per_group = 4
    176.     return ResNet(Bottleneck, [3, 4, 6, 3],
    177.                   num_classes=num_classes,
    178.                   include_top=include_top,
    179.                   groups=groups,
    180.                   width_per_group=width_per_group)
    181.  
    182.  
    183. def resnext101_32x8d(num_classes=1000, include_top=True):
    184.     # https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth
    185.     groups = 32
    186.     width_per_group = 8
    187.     return ResNet(Bottleneck, [3, 4, 23, 3],
    188.                   num_classes=num_classes,
    189.                   include_top=include_top,
    190.                   groups=groups,
    191.                   width_per_group=width_per_group)
    192.  
    193. if __name__ == "__main__":
    194.     resnet = ResNet(BasicBlock, [3, 4, 6, 3], num_classes=5)
    195.     in_data = torch.randn(1, 3, 224, 224)
    196.     out = resnet(in_data)
    197.     print(out)

    完成网络的定义之后，可以单独执行一下这个文件，用来验证网络定义的是否正确。如果可以正确输出，就没问题。

    在这里输出为

    tensor([[-0.4490,  0.5792, -0.5026, -0.6024,  0.1399]],
    grad_fn=)

    说明网络定义正确。

    三.开始训练

     加载数据集

    首先定义一个字典，用于用于对train和val进行预处理，包括裁剪成224*224大小，训练集随机水平翻转（一般验证集不需要此操作），转换成张量，图像归一化。

    然后利用DataLoader模块加载数据集，并设置batch_size为16，同时，设置数据加载器的工作进程数nw，加快速度。

    1. import os
    2. import sys
    3. import json
    4.  
    5. import torch
    6. import torch.nn as nn
    7. import torch.optim as optim
    8. from torchvision import transforms, datasets
    9. from torch.utils.data import DataLoader
    10. from tqdm import tqdm
    11.  
    12. from model import resnet34
    13.  
    14.  
    15. def main():
    16.     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    17.     print(f"using {device} device.")
    18.  
    19.     data_transform = {
    20.         "train": transforms.Compose([transforms.RandomResizedCrop(224),
    21.                                      transforms.RandomHorizontalFlip(),
    22.                                      transforms.ToTensor(),
    23.                                      transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
    24.         "val": transforms.Compose([transforms.Resize(256),
    25.                                    transforms.CenterCrop(224),
    26.                                    transforms.ToTensor(),
    27.                                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}
    28.  
    29.     # 获取数据集路径
    30.     image_path = os.path.join(os.getcwd(), "data_set", "flower_data")
    31.     assert os.path.exists(image_path), f"{image_path} path does not exist."
    32.     # 加载数据集，准备读取
    33.     train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"), transform=data_transform["train"])
    34.     validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"), transform=data_transform["val"])
    35.  
    36.     nw = min([os.cpu_count(), 16 if 16 > 1 else 0, 8])  # number of workers，加速图像预处理
    37.     print(f'Using {nw} dataloader workers every process')
    38.  
    39.     # 加载数据集
    40.     train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=nw)
    41.     validate_loader = torch.utils.data.DataLoader(validate_dataset, batch_size=16, shuffle=False, num_workers=nw)
    42.     train_num = len(train_dataset)
    43.     val_num = len(validate_dataset)
    44.     print(f"using {train_num} images for training, {val_num} images for validation.")

    生成json文件

    将训练数据集的类别标签转换为字典格式，并将其写入名为'class_indices.json'的文件中。

    1. 从train_dataset中获取类别标签到索引的映射关系，存储在flower_list变量中。
    2. 使用列表推导式将flower_list中的键值对反转，得到一个新的字典cla_dict，其中键是原始类别标签，值是对应的索引。
    3. 使用json.dumps()函数将cla_dict转换为JSON格式的字符串，设置缩进为4个空格。
    4. 使用with open()语句以写入模式打开名为'class_indices.json'的文件，并将JSON字符串写入文件
    1.    # {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4} 雏菊 蒲公英 玫瑰 向日葵 郁金香
    2.     # 从训练集中获取类别标签到索引的映射关系，存储在flower_list变量
    3.     flower_list = train_dataset.class_to_idx
    4.     # 使用列表推导式将flower_list中的键值对反转，得到一个新的字典cla_dict
    5.     cla_dict = dict((val, key) for key, val in flower_list.items())
    6.     # write dict into json file
    7.     json_str = json.dumps(cla_dict, indent=4)
    8.     with open('class_indices.json', 'w') as json_file:
    9.         json_file.write(json_str)

    加载预训练模型开始训练

    首先定义网络对象net，在这里我们使用了迁移学习来使网络训练效果更好；使用net.fc = nn.Linear(in_channel, 5)设置输出类别数（这里为5）；训练10轮，并使用train_bar = tqdm(train_loader, file=sys.stdout)来可视化训练进度条，之后再进行反向传播和参数更新;同时，每一轮训练完成都要进行学习率更新；之后开始对验证集进行计算精确度，完成后保存模型。

    1.     # load pretrain weights
    2.     # download url: https://download.pytorch.org/models/resnet34-333f7ec4.pth
    3.     net = resnet34()
    4.     model_weight_path = "./resnet34-pre.pth"
    5.     assert os.path.exists(model_weight_path), f"file {model_weight_path} does not exist."
    6.     net.load_state_dict(torch.load(model_weight_path, map_location='cpu'))
    7.  
    8.     # change fc layer structure
    9.     in_channel = net.fc.in_features
    10.     net.fc = nn.Linear(in_channel, 5)
    11.     net.to(device)
    12.     loss_function = nn.CrossEntropyLoss()
    13.     optimizer = optim.Adam([p for p in net.parameters() if p.requires_grad], lr=0.0001)
    14.  
    15.     epochs = 10
    16.     best_acc = 0.0
    17.     train_steps = len(train_loader)
    18.     for epoch in range(epochs):
    19.         # train
    20.         net.train()
    21.         running_loss = 0.0
    22.         train_bar = tqdm(train_loader, file=sys.stdout)
    23.         for step, data in enumerate(train_bar):
    24.             images, labels = data
    25.             optimizer.zero_grad()
    26.             logits = net(images.to(device))
    27.             loss = loss_function(logits, labels.to(device))
    28.             loss.backward()
    29.             optimizer.step()
    30.  
    31.             # print statistics
    32.             running_loss += loss.item()
    33.  
    34.             train_bar.desc = "train epoch[{epoch + 1}/{epochs}] loss:{loss:.3f}"
    35.  
    36.         # validate
    37.         net.eval()
    38.         acc = 0.0  # accumulate accurate number / epoch
    39.         with torch.no_grad():
    40.             val_bar = tqdm(validate_loader, file=sys.stdout)
    41.             for val_data in val_bar:
    42.                 val_images, val_labels = val_data
    43.                 outputs = net(val_images.to(device))
    44.                 # loss = loss_function(outputs, test_labels)
    45.                 predict_y = torch.max(outputs, dim=1)[1]
    46.                 acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
    47.  
    48.                 val_bar.desc = f"valid epoch[{epoch + 1}/{epochs}]"
    49.  
    50.         val_accurate = acc / val_num
    51.         print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
    52.               (epoch + 1, running_loss / train_steps, val_accurate))
    53.  
    54.         if val_accurate > best_acc:
    55.             best_acc = val_accurate
    56.             torch.save(net, "./resnet.pth")
    57.  
    58.     print('Finished Training')

    最后对代码进行整理，完整的train.py如下

    1. import os
    2. import sys
    3. import json
    4.  
    5. import torch
    6. import torch.nn as nn
    7. import torch.optim as optim
    8. from torchvision import transforms, datasets
    9. from torch.utils.data import DataLoader
    10. from tqdm import tqdm
    11.  
    12. from model import resnet34
    13.  
    14.  
    15. def main():
    16.     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    17.     print(f"using {device} device.")
    18.  
    19.     data_transform = {
    20.         "train": transforms.Compose([transforms.RandomResizedCrop(224),
    21.                                      transforms.RandomHorizontalFlip(),
    22.                                      transforms.ToTensor(),
    23.                                      transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
    24.         "val": transforms.Compose([transforms.Resize(256),
    25.                                    transforms.CenterCrop(224),
    26.                                    transforms.ToTensor(),
    27.                                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}
    28.  
    29.     # 获取数据集路径
    30.     image_path = os.path.join(os.getcwd(), "data_set", "flower_data")
    31.     assert os.path.exists(image_path), f"{image_path} path does not exist."
    32.     # 加载数据集，准备读取
    33.     train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"), transform=data_transform["train"])
    34.     validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"), transform=data_transform["val"])
    35.  
    36.     nw = min([os.cpu_count(), 16 if 16 > 1 else 0, 8])  # number of workers，加速图像预处理
    37.     print(f'Using {nw} dataloader workers every process')
    38.  
    39.     # 加载数据集
    40.     train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=nw)
    41.     validate_loader = torch.utils.data.DataLoader(validate_dataset, batch_size=16, shuffle=False, num_workers=nw)
    42.     train_num = len(train_dataset)
    43.     val_num = len(validate_dataset)
    44.     print(f"using {train_num} images for training, {val_num} images for validation.")
    45.  
    46.     # {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4} 雏菊 蒲公英 玫瑰 向日葵 郁金香
    47.     # 从训练集中获取类别标签到索引的映射关系，存储在flower_list变量
    48.     flower_list = train_dataset.class_to_idx
    49.     # 使用列表推导式将flower_list中的键值对反转，得到一个新的字典cla_dict
    50.     cla_dict = dict((val, key) for key, val in flower_list.items())
    51.     # write dict into json file
    52.     json_str = json.dumps(cla_dict, indent=4)
    53.     with open('class_indices.json', 'w') as json_file:
    54.         json_file.write(json_str)
    55.     
    56.  
    57.     # load pretrain weights
    58.     # download url: https://download.pytorch.org/models/resnet34-333f7ec4.pth
    59.     net = resnet34()
    60.     model_weight_path = "./resnet34-pre.pth"
    61.     assert os.path.exists(model_weight_path), f"file {model_weight_path} does not exist."
    62.     net.load_state_dict(torch.load(model_weight_path, map_location='cpu'))
    63.  
    64.     # change fc layer structure
    65.     in_channel = net.fc.in_features
    66.     net.fc = nn.Linear(in_channel, 5)
    67.     net.to(device)
    68.     loss_function = nn.CrossEntropyLoss()
    69.     optimizer = optim.Adam([p for p in net.parameters() if p.requires_grad], lr=0.0001)
    70.  
    71.     epochs = 10
    72.     best_acc = 0.0
    73.     train_steps = len(train_loader)
    74.     for epoch in range(epochs):
    75.         # train
    76.         net.train()
    77.         running_loss = 0.0
    78.         train_bar = tqdm(train_loader, file=sys.stdout)
    79.         for step, data in enumerate(train_bar):
    80.             images, labels = data
    81.             optimizer.zero_grad()
    82.             logits = net(images.to(device))
    83.             loss = loss_function(logits, labels.to(device))
    84.             loss.backward()
    85.             optimizer.step()
    86.  
    87.             # print statistics
    88.             running_loss += loss.item()
    89.  
    90.             train_bar.desc = "train epoch[{epoch + 1}/{epochs}] loss:{loss:.3f}"
    91.  
    92.         # validate
    93.         net.eval()
    94.         acc = 0.0  # accumulate accurate number / epoch
    95.         with torch.no_grad():
    96.             val_bar = tqdm(validate_loader, file=sys.stdout)
    97.             for val_data in val_bar:
    98.                 val_images, val_labels = val_data
    99.                 outputs = net(val_images.to(device))
    100.                 # loss = loss_function(outputs, test_labels)
    101.                 predict_y = torch.max(outputs, dim=1)[1]
    102.                 acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
    103.  
    104.                 val_bar.desc = f"valid epoch[{epoch + 1}/{epochs}]"
    105.  
    106.         val_accurate = acc / val_num
    107.         print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
    108.               (epoch + 1, running_loss / train_steps, val_accurate))
    109.  
    110.         if val_accurate > best_acc:
    111.             best_acc = val_accurate
    112.             torch.save(net, "./resnet.pth")
    113.  
    114.     print('Finished Training')
    115.  
    116.  
    117. if __name__ == '__main__':
    118.     main()

    四.模型预测

    新建一个predict.py文件用于预测，将输入图像处理后转换成张量格式，img = torch.unsqueeze(img, dim=0)是在输入图像张量 img 的第一个维度上增加一个大小为1的维度，因此将图像张量的形状从 [通道数， 高度， 宽度 ] 转换为 [1, 通道数， 高度， 宽度]。然后加载模型进行预测，并打印出结果，同时可视化。

    1. import os
    2. import json
    3.  
    4. import torch
    5. from PIL import Image
    6. from torchvision import transforms
    7. import matplotlib.pyplot as plt
    8.  
    9. from model import resnet34
    10.  
    11.  
    12. def main():
    13.     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    14.  
    15.     data_transform = transforms.Compose(
    16.         [transforms.Resize(256),
    17.          transforms.CenterCrop(224),
    18.          transforms.ToTensor(),
    19.          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    20.  
    21.     # load image
    22.     img = Image.open("./2536282942_b5ca27577e.jpg")
    23.     plt.imshow(img)
    24.     # [N, C, H, W]
    25.     img = data_transform(img)
    26.     # expand batch dimension
    27.     # 在输入图像张量 img 的第一个维度上增加一个大小为1的维度
    28.     # 将图像张量的形状从 [通道数， 高度， 宽度 ] 转换为 [1, 通道数， 高度， 宽度]
    29.     img = torch.unsqueeze(img, dim=0)
    30.  
    31.     # read class_indict
    32.     with open('./class_indices.json', "r") as f:
    33.         class_indict = json.load(f)
    34.  
    35.     # create model
    36.     model = resnet34(num_classes=5).to(device)
    37.     model = torch.load("./resnet34.pth")
    38.  
    39.     # prediction
    40.     model.eval()
    41.     with torch.no_grad():
    42.         # predict class
    43.         output = torch.squeeze(model(img.to(device))).cpu()
    44.         predict = torch.softmax(output, dim=0)
    45.         predict_class = torch.argmax(predict).numpy()
    46.  
    47.     print_result = f"class: {class_indict[str(predict_class)]}   prob: {predict[predict_class].numpy():.3}"
    48.     plt.title(print_result)
    49.     for i in range(len(predict)):
    50.         print(f"class: {class_indict[str(i)]:10}   prob: {predict[i].numpy():.3}")
    51.     plt.show()
    52.  
    53.  
    54. if __name__ == '__main__':
    55.     main()

    预测结果

    五.模型可视化

    将生成的pth文件导入netron工具，可视化结果为

    发现很不清晰，因此将它转换成多用于嵌入式设备部署的onnx格式

    编写onnx.py

    1. import torch
    2. import torchvision
    3. from model import resnet34
    4.  
    5. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    6. model = resnet34(num_classes=5).to(device)
    7. model=torch.load("/home/lm/Resnet/resnet34.pth")
    8. model.eval()
    9. example = torch.ones(1, 3, 244, 244)
    10. example = example.to(device)
    11. torch.onnx.export(model, example, "resnet34.onnx", verbose=True, opset_version=11)

      将生成的onnx文件导入，这样的可视化清晰了许多

     六.批量数据预测

    现在新建一个dta文件夹，里面放入五类带预测的样本，编写代码完成对整个文件夹下所有样本的预测，即批量预测。

    batch_predict.py如下：

    1. import os
    2. import json
    3.  
    4. import torch
    5. from PIL import Image
    6. from torchvision import transforms
    7.  
    8. from model import resnet34
    9.  
    10.  
    11. def main():
    12.     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    13.  
    14.     data_transform = transforms.Compose(
    15.         [transforms.Resize(256),
    16.          transforms.CenterCrop(224),
    17.          transforms.ToTensor(),
    18.          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    19.  
    20.     # load image
    21.     # 指向需要遍历预测的图像文件夹
    22.     imgs_root = "./data/imgs"
    23.     # 读取指定文件夹下所有jpg图像路径
    24.     img_path_list = [os.path.join(imgs_root, i) for i in os.listdir(imgs_root) if i.endswith(".jpg")]
    25.  
    26.     # read class_indict
    27.     json_file = open('./class_indices.json', "r")
    28.     class_indict = json.load(json_file)
    29.  
    30.     # create model
    31.     model = resnet34(num_classes=5).to(device)
    32.     model = torch.load("./resnet34.pth")
    33.  
    34.     # prediction
    35.     model.eval()
    36.     batch_size = 8  # 每次预测时将多少张图片打包成一个batch
    37.     with torch.no_grad():
    38.         for ids in range(0, len(img_path_list) // batch_size):
    39.             img_list = []
    40.             for img_path in img_path_list[ids * batch_size: (ids + 1) * batch_size]:
    41.                 img = Image.open(img_path)
    42.                 img = data_transform(img)
    43.                 img_list.append(img)
    44.  
    45.             # batch img
    46.             # 将img_list列表中的所有图像打包成一个batch
    47.             batch_img = torch.stack(img_list, dim=0)
    48.             # predict class
    49.             output = model(batch_img.to(device)).cpu()
    50.             predict = torch.softmax(output, dim=1)
    51.             probs, classes = torch.max(predict, dim=1)
    52.  
    53.             for idx, (pro, cla) in enumerate(zip(probs, classes)):
    54.                 print(f"image: {img_path_list[ids*batch_size+idx]}  class: {class_indict[str(cla.numpy())]}  prob: {pro.numpy():.3}")
    55.  
    56.  
    57. if __name__ == '__main__':
    58.     main()

    运行之后，输出

    image: ./data/imgs/455728598_c5f3e7fc71_m.jpg  class: dandelion  prob: 0.989
    image: ./data/imgs/3464015936_6845f46f64.jpg  class: dandelion  prob: 0.999
    image: ./data/imgs/3461986955_29a1abc621.jpg  class: dandelion  prob: 0.996
    image: ./data/imgs/8223949_2928d3f6f6_n.jpg  class: dandelion  prob: 0.991
    image: ./data/imgs/10919961_0af657c4e8.jpg  class: dandelion  prob: 1.0
    image: ./data/imgs/10443973_aeb97513fc_m.jpg  class: dandelion  prob: 0.906
    image: ./data/imgs/8475758_4c861ab268_m.jpg  class: dandelion  prob: 0.805
    image: ./data/imgs/3857059749_fe8ca621a9.jpg  class: dandelion  prob: 1.0
    image: ./data/imgs/2457473644_5242844e52_m.jpg  class: dandelion  prob: 1.0
    image: ./data/imgs/146023167_f905574d97_m.jpg  class: dandelion  prob: 0.998
    image: ./data/imgs/2502627784_4486978bcf.jpg  class: dandelion  prob: 0.488
    image: ./data/imgs/2481428401_bed64dd043.jpg  class: dandelion  prob: 1.0
    image: ./data/imgs/13920113_f03e867ea7_m.jpg  class: dandelion  prob: 1.0
    image: ./data/imgs/2535769822_513be6bbe9.jpg  class: dandelion  prob: 0.997
    image: ./data/imgs/3954167682_128398bf79_m.jpg  class: dandelion  prob: 1.0
    image: ./data/imgs/2516714633_87f28f0314.jpg  class: dandelion  prob: 0.998
    image: ./data/imgs/2634665077_597910235f_m.jpg  class: dandelion  prob: 0.996
    image: ./data/imgs/3502447188_ab4a5055ac_m.jpg  class: dandelion  prob: 0.999
    image: ./data/imgs/425800274_27dba84fac_n.jpg  class: dandelion  prob: 0.422
    image: ./data/imgs/3365850019_8158a161a8_n.jpg  class: dandelion  prob: 1.0
    image: ./data/imgs/674407101_57676c40fb.jpg  class: dandelion  prob: 1.0
    image: ./data/imgs/2628514700_b6d5325797_n.jpg  class: dandelion  prob: 0.999
    image: ./data/imgs/3688128868_031e7b53e1_n.jpg  class: dandelion  prob: 0.962
    image: ./data/imgs/2502613166_2c231b47cb_n.jpg  class: dandelion  prob: 1.0

    完成预期功能（这里我的样本都是dandelion,当然混合的也可以）

    七.模型改进

    当不加载预训练模型，而从头开始训练的话，当epoch为50时，经实际训练，准确率为80%多，但当加载预训练模型时，完成第一次迭代准确率就已达到了90%，这也正说明了迁移学习的好处。

    同时，这里采用的是Resnet34，也可以尝试更深的50、101、152层网络。

    还有其他方法会在之后进行补充。

  • 相关阅读:
    舵机调试上位机
    猿创征文｜2022 互联网从业心得：“但愿人长久，千里共婵娟”，中秋先学会与自己内心团圆
    JSP response,request操作中(中文乱码)-如何解决呢？
    PC微信3.9.7内测版，更新功能一览（附下载）
    使用R语言进行简单的主成分分析（PCA）
    Java_断点调试/类变量/类方法/main语法
    vue中预览docx
    MySQL-事务隔离-2
    深拷贝与浅拷贝
    如何在Ubuntu中挂载新硬盘
  • 原文地址：https://blog.csdn.net/qq_46454669/article/details/133977165