基于Resnet的钢材表面缺陷分类问题

ccccccwen

已于 2022-08-25 09:37:26 修改

阅读量1.1k

点赞数 1

文章标签：分类 python 机器学习深度学习

于 2022-08-24 22:06:43 首次发布

本文链接：https://blog.csdn.net/weixin_62064869/article/details/126511604

版权

1.1 创建对应类别文件夹并且进行同类分类

1.2 将训练数据按比例分开训练集和验证集

前言

本文是主要讲述在灰度图如何进行迁移学习的图像分类任务，采用的是钢材表面缺陷数据集，并且使用Resnet迁移学习进行训练。

一、数据集

本文采用的是钢材表面缺陷数据集，该数据集的特点有：

该数据集主要分有七类，在训练中分别用0到6代替其中的类型
该数据集全都是单通道的灰度图，图片后缀名为.bmp
该数据集经过分类训练集有2000张照片，测试集有400张
该数据集有training.csv，包括了训练集的id和label

数据集类别如图所示，从左到右分别为0、1、2、3、4、5、6类别：

training.csv部分内容如图所示：

二、数据预处理

1.将图片随机分到不同的文件夹

1.1 创建对应类别文件夹并且进行同类分类

- 创建一个TrainAll文件夹，其中包含0、1、2、3、4、5、6七个文件夹

- 输入混合的训练集图片路径，将混合的照片按照training.csv分到对应标签的文件夹

def integration():
    """根据csv文件构建文件夹"""
    #读取csv文件，计算长度
    data_x= pd.read_csv(filepath_or_buffer = 'training.csv', sep = ',')["Image_name	/(.bmp)"].values
    data_y= pd.read_csv(filepath_or_buffer = 'training.csv', sep = ',')["categories"].values
    num = len(data_x)
    #如果原路径存在该文件，先清空了，保证数据准确
    if os.path.exists('TrainAll') :
        shutil.rmtree('TrainAll')
    #创建训练分类文件夹
    for i in range(8):
        i = str(i)
        os.makedirs(os.path.join('TrainAll',i))
    #输入混合图片路径
    dir = input('请输入训练集的文件夹的路径：')
    root = ''
    for _root, _dirs, _files in os.walk(dir):  # root为输入的文件夹名，dirs为子文件夹的名字，files为图片
        root = _root
        print("训练集上有{}张照片".format(len(_files)))
    #对训练集上的数据进行分类
    for j in range(8):
        label = str(j)
        for i in range(num) :
            if data_y[i] == j:
                i = ''+ str(i+1)+'.bmp'
                shutil.copy(os.path.join(root,i), os.path.join('TrainAll',label))

1.2 将训练数据按比例分开训练集和验证集

- 因为该数据集的测试集未公布，所以训练的过程中需要在训练集中划分出验证集

- 输入训练集和验证集的比例，代码会自动按照比例进行分配

def Classify ():
    """按比例分类训练集和测试集"""
    dir = 'TrainAll'
    proportion = input ('请输入训练集占全集（验证集和训练集总和）的比例（小数）：') #输入比例
    proportion  = float(proportion)
    if os.path.exists('data'):
        shutil.rmtree('data')  # 删除原文件夹,保证数据准确
    print("正在分类")
    for _root,_dirs,_files in os.walk(dir):     # root为输入的文件夹名，dirs为子文件夹的名字，files为图片
        for _name in _dirs:                    # 遍历
            folder = os.path.join(_root,_name) # 将文件名和子文件名合成一个路径
            moveFile(folder,_name,proportion)     #划分新的数据集
        print("分类完成！")
        break #循环一次就行

2.加载数据集和验证集

- 训练集和验证集分别进行不同的图像增强，分别进行加载数据

- 必须有transforms.Grayscale(1)读入灰度图，否则为RGB

- transforms.Normalize中的0.485, 0.229是在训练集中灰度图中计算得到的均值和方差

    # 加载数据
    data_transform = {
        "train": transforms.Compose([transforms.Grayscale(1)
                                     transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize(0.485, 0.229, inplace=True)]),
        "val": transforms.Compose([transforms.Grayscale(1)
                                   transforms.Resize(256),
                                   transforms.CenterCrop(224),
                                   transforms.ToTensor(),
                                   transforms.Normalize(0.485, 0.229, inplace=True)])}
    image_path = os.path.join("data")
    train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
                                         transform=data_transform["train"])
    train_num = len(train_dataset)
    # {'0':0, '1':1, '2':2, '3':3, '4':4, '5':5, '6':6, '7':7}
    list = train_dataset.class_to_idx
    cla_dict = dict((val, key) for key, val in list.items())
    # 写json文件，方便对应类别
    json_str = json.dumps(cla_dict, indent=4)
    with open('class_indices.json', 'w') as json_file:
        json_file.write(json_str)
    batch_size = 8
    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batch_size, shuffle=True,
                                               num_workers=0)

    validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
                                            transform=data_transform["val"])
    val_num = len(validate_dataset)
    validate_loader = torch.utils.data.DataLoader(validate_dataset,
                                                  batch_size=batch_size, shuffle=False,
                                                  num_workers=0)
    print("训练集有{}张照片, 测试集有{}张照片".format(train_num,val_num))

二.模型加载

1.模型选择

- 在保证准确率的同时，提高分类速度，选择了resnet-18的模型结构

- 因为分类任务类似，因此使用resnet-18 提供的预训练权重进行训练

2.模型搭建

- 因为是灰度图，所以class resnet中的self.conv1中的输入通道应该为1

class BasicBlock(nn.Module):
    expansion = 1
    def __init__(self, in_channel, out_channel, stride=1, downsample=None, **kwargs):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,
                                kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,
                               kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)
        self.downsample = downsample

    def forward(self, x):
        identity = x
        if self.downsample is not None:
            identity = self.downsample(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out += identity
        out = self.relu(out)
        return out


class ResNet(nn.Module):
    def __init__(self,
                 block,
                 blocks_num,
                 num_classes=1000,
                 include_top=True,
                 groups=1,
                 width_per_group=64):
        super(ResNet, self).__init__()
        self.include_top = include_top
        self.in_channel = 64
        self.groups = groups
        self.width_per_group = width_per_group
        self.conv1 = nn.Conv2d(1, self.in_channel, kernel_size=7, stride=2,
                               padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(self.in_channel)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, blocks_num[0])
        self.layer2 = self._make_layer(block, 128, blocks_num[1], stride=2)
        self.layer3 = self._make_layer(block, 256, blocks_num[2], stride=2)
        self.layer4 = self._make_layer(block, 512, blocks_num[3], stride=2)
        if self.include_top:
            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # output size = (1, 1)
            self.fc = nn.Linear(512 * block.expansion, num_classes)
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')

    def _make_layer(self, block, channel, block_num, stride=1):
        downsample = None
        if stride != 1 or self.in_channel != channel * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(channel * block.expansion))
        layers = []
        layers.append(block(self.in_channel,
                            channel,
                            downsample=downsample,
                            stride=stride,
                            groups=self.groups,
                            width_per_group=self.width_per_group))
        self.in_channel = channel * block.expansion
        for _ in range(1, block_num):
            layers.append(block(self.in_channel,
                                channel,
                                groups=self.groups,
                                width_per_group=self.width_per_group))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        if self.include_top:
            x = self.avgpool(x)
            x = torch.flatten(x, 1)
            x = self.fc(x)
        return x

def resnet18(num_classes=1000, include_top=True):
    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes, include_top=include_top)

3.实例化模型

- 将下载好的预训练权重放到与代码同以路径下，读入权重

- 将全连接层1000改为8

    #加载模型
    net = resnet18()
    model_weight_path = "./resnet18-5c106cde.pth"
    pre_state_dict = torch.load(model_weight_path)
    new_state_dict = {}
    #遍历修改模型的各个层
    for k, v in net.state_dict().items():
        # 如果原模型的层也在新模型的层里面， 那新模型就加载原先训练好的权重
        if k in pre_state_dict.keys() and k != 'conv1.weight':
            new_state_dict[k] = pre_state_dict[k]  
    net.load_state_dict(new_state_dict, False)
    in_channel = net.fc.in_features
    net.fc = nn.Linear(in_channel, 8)
    net.to(device)

三.损失函数

- 使用交叉熵损失函数

- 使用Afam学习率为0.0001的优化方法

    #损失函数
    loss_function = nn.CrossEntropyLoss()
    params = [p for p in net.parameters() if p.requires_grad]
    optimizer = optim.Adam(params, lr=0.0001)

四.训练

- 运用了tqdm进行可视化进程

- 方向梯度更新权重

- 比较验证集的准确率，保存准确率高的权重

#训练
    for epoch in range(epochs):
        net.train()
        running_loss = 0.0
        train_bar = tqdm(train_loader, file=sys.stdout)
        for step, data in enumerate(train_bar):
            images, labels = data
            optimizer.zero_grad()
            logits = net(images.to(device))
            loss = loss_function(logits, labels.to(device))
            loss.backward()
            optimizer.step()
            #可视化数据
            running_loss += loss.item()
            train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,epochs, loss)
        #检测
        net.eval()
        acc = 0.0
        with torch.no_grad():
            val_bar = tqdm(validate_loader, file=sys.stdout)
            for val_data in val_bar:
                val_images, val_labels = val_data
                outputs = net(val_images.to(device))
                predict_y = torch.max(outputs, dim=1)[1]
                acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
                val_bar.desc = "valid epoch[{}/{}]".format(epoch + 1,
                                                           epochs)
        val_accurate = acc / val_num
        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
              (epoch + 1, running_loss / train_steps, val_accurate))
        #只保留最好的预测模型
        if val_accurate > best_acc:
            best_acc = val_accurate
            torch.save(net.state_dict(), save_path)
    print('训练完成！')
    print('训练模型已保存到当前目录')