深度学习中，分离分类错误图像的方法（本人用于GAN数据集扩充）

省流：

事先获取图像标签集作为列表，并创建对应的文件夹。在检测验证集时，获取分类失误的图像标签。最终将图像保存到对应的文件夹中。

---

详细：

目的：

用GAN扩充数据集是为了解决，数据集部分类型数据较少的问题。在用GAN扩充数据集时，我们不能盲目扩充数据，否则会遇到数据分布不均衡、过拟合的情况，导致分类效果下降。这需要我们事先调查哪些图像容易被模型误判。

具体步骤：

步骤一：创建验证集与错题集地址

由于神经网络模型输出的标签是数字的形式，所以需要对当前标签对照文件排列顺序创建列表，这里以WM0-811K晶圆数据集为例

创建验证集与错题集地址

label_name = ['Center', 'Donut', 'Edge-Loc', 'Edge-Ring', 'Loc', 'Near-full', 'none', 'Random', 'Scratch']

test_dataset = CustomDataset(transform=transform, train=False)

    # 创建数据加载器实例
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

# 创建多个列表用于存储准确率、召回率、F1参数
cor = [0] * 9        # 存储各类别 检测正确的数量
lab = [0] * 9        # 存储各类别 总的数量
pre = [0] * 9        # 存储各类别 被预测到的数量（并不要求预测对）
acc = [0] * 9        # 存储各类别 准确率
rec = [0] * 9        # 存储各类别 召回率
f1 = [0] * 9         # 存储各类别 F1参数
tot = [0] * 3        # 存储总 准确率、召回率、F1参数

# 创建错题集地址
error_folder = 'E:/***********/GAN/DCGAN/ResNet_pt/errorset'
for label in label_name:
    creat_path(os.path.join(error_folder,t,label))

步骤二：检测验证集

加载训练好的模型，将验证集输入模型进行检测

model = torch.load(org_path)        # 加载训练模型
num = 0
for i, (images, labels) in enumerate(test_dataloader):        # 逐批读取验证集

逐一检测验证集，取输出值中最大值的索引作为样本标签。同时统计验证集的各个类别样本数量，也就是各个标签准确率的分母，用于后期准确率的计算。

    # print(type(labels),labels.shape)        # <class 'torch.Tensor'> torch.Size([1])
    # print(labels.shape)
    for l in range(len(labels)):
        lab[labels[l]] += 1     # 统计实际的样本
    images = images.to(device)
    labels = labels.to(device)
    logits, probas = model(images)      #
    _, pred_labels = torch.max(probas, 1)      # 选择probas行中最大值的索引作为标签

步骤三：样本统计与保存

统计预测的各类别被预测到的样本、预测正确的样本分离误判样本并保存

    dim = images.shape

    for j in range(len(pred_labels)):
        pre[pred_labels[j]] += 1     # 统计预测的样本
        # print(labels[j],pred_labels[j])
        if labels[j] == pred_labels[j]:
            cor[labels[j]] += 1     # 统计预测正确的样本情况
        else:        # 将标签与样本从GPU中转移到CPU中并数组化
            k = labels[j].cpu().numpy()
            # print(label_name[k])
            image = images[j,:,:,:].reshape(dim[1],dim[2],dim[3]).cpu().numpy()
            img_path = os.path.join(error_folder,t,label_name[k],f'error_{num}.png')
            cv2.imwrite(img_path,np.transpose(image*255, (1, 2, 0)))
         # 将数据格式从 H,W,C 转为 C,H,W 模式，用于cv2.imwrite函数保存
            num += 1

for i in range(len(cor)):
    acc[i] = cor[i]/lab[i]      # 计算准确率
    try:
        rec[i] = cor[i]/pre[i]      # 计算召回率
        f1[i] = 2 * acc[i] * rec[i]/ (acc[i] + rec[i])
    except:
        print(i)
    tot[0] += lab[i]
    tot[1] += cor[i]
    # print('accuracy of ' + label_name[i] + ': {:.5f}' .format(acc[i]))

tot[2] = tot[1]/tot[0]
print(tot[2])

如有谬误请各位支出，谢谢！

---

完整代码

import cv2
import time
import os

import json
import numpy as np
import random
import torch
import torch.nn.functional as F
import torch.nn as nn

from sklearn.metrics import accuracy_score,precision_score,recall_score,f1_score

from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import hiddenlayer as hl

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def creat_path(img_Topath):
    if not os.path.exists(img_Topath):
        os.makedirs(img_Topath)
        
# 对标签进行编码
label_encoder = {'Center': 0, 'Donut': 1, 'Edge-Loc': 2, 'Edge-Ring': 3, 'Loc': 4, 'Near-full': 5, 'none': 6,
                 'Random': 7, 'Scratch': 8, '_Org': 9}
label_name = ['Center', 'Donut', 'Edge-Loc', 'Edge-Ring', 'Loc', 'Near-full', 'none', 'Random', 'Scratch']

创建json数据库存储、划分样本的地址

org_path = f"E:/************/GAN/DCGAN/ResNet_pt/{t}.pt"       # 模型的保存地址
org_data = 'E:/************'        # 样本的存储地址

def generate_datasets(root = org_data ):      # 原图
    """
    对数据集进行划分  每个标签按照28比例进行划分
    :param root: 数据所在目录
    :return:
    """
    # 遍历所有的目录
    dir_names = [i for i in os.listdir(root) if os.path.isdir(os.path.join(root, i))]
    # 保存数据
    data = {"train": {}, "val": {}}
    ratio = 0.2
    for dir_name in dir_names:
        # 对所有的标签目录下进行遍历
        train = []
        test = []
        # 获取所有的文件名称
        file_names = os.listdir(os.path.join(root, dir_name))
        random.shuffle(file_names)

        # 获取划分索引

        
        # 如果是_Org 或者none类型的，只用其中3000张图像
        if len(file_names) > 3000:
            file_names = file_names[:3000]
        index_split = int(len(file_names) * ratio)
        # 根据索引进行数据集划分
        for index, file_name in enumerate(file_names):
            if '76-' in file_name:                  # 将生成的图加入
                train.append(os.path.join(root, dir_name, file_name))
            elif len(test) < index_split:               # 如果测试集没有装满，就装入测试集
                test.append(os.path.join(root, dir_name, file_name))
            else:
                train.append(os.path.join(root, dir_name, file_name))

        # 将数据保存到data中
        data["train"][dir_name] = train
        data["val"][dir_name] = test
        
    # 将json数据保存到本地
    with open("org.json", "w") as fp:
        json.dump(data, fp)

加载数据集与模型

def load_json(data_path: str = "data.json"):
    """
    加载json数据
    :param data_path:文件路径
    :return:
    """
    with open(data_path, "r") as fp:
        data = json.load(fp)
    return data

class CustomDataset(Dataset):
    # 自定义数据集
    def __init__(self, transform=None, train=True):
        self.transform = transform
        data = load_json("org.json")       # 读取data.json文件中所有的数据地址
        if train:
            data = data["train"]        # 选择要加载的部分
        else:
            data = data["val"]

        self.image_paths = []       # 地址集
        self.labels = []        # 标签集

        for label, value in data.items():       # 这里data是个字典
            # pdb.set_trace()
            temp = []       # 保存一类标签标签的 标签集
            for i in range(len(value)):     # 确认某类的数据数量
                temp.append(label_encoder[label])       # 将标签数字化，传入temp：  [0,0,0,0,0,...,0]               
            self.labels.extend(temp)        # 将标签集temp装载入self.labels
            # self.labels.extend([label_encoder[label] for _ in range(len(value))])
            self.image_paths.extend(value)      # 将地址装入image_paths

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        image = Image.open(image_path).convert('RGB')       # 将image_path对应的图读入image
        if self.transform:
            image = self.transform(image)
        return image, self.labels[idx]
'''
1x1 , 2^n
3x3 , 2^n
1x1 , 2^(n+2)
'''
class Bottleneck(nn.Module):        # 构建残差块，继承了torch.nn.Module,重写了__init__和forward
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):    # 
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)     # 输入通道数、输出通道数、卷积核、偏置
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,     # 输入通道数、输出通道数、卷积核、偏置
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out




class ResNet(nn.Module):
                # Bottleneck,[3, 4, 6, 3],类的总数,是否黑白化
    def __init__(self, block, layers, num_classes, grayscale):
        self.inplanes = 64
        if grayscale:
            in_dim = 1
        else:
            in_dim = 3
        super(ResNet, self).__init__()
        self.conv1 = nn.Conv2d(in_dim, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        self.fc = nn.Linear(4608 * block.expansion, num_classes)        # 改成 4608 (9*512)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, (2. / n)**.5)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        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)
        # because MNIST is already 1x1 here:
        # disable avg pooling
        #x = self.avgpool(x)
        
        x = x.view(x.size(0), -1)
        logits = self.fc(x)
        probas = F.softmax(logits, dim=1)
        return logits, probas



def resnet50(n_classes):
    """Constructs a ResNet-34 model."""
    model = ResNet(block=Bottleneck, 
                   layers=[3, 4, 6, 3],
                   num_classes=n_classes,
                   grayscale=False)     # 
    return model

batch_size = 32 # 批次大小
num_classes = 9  # 根据你的数据集调整类别数
epochs = 15  # 训练轮数
learning_rate = 0.001 # 学习率

history = hl.History()
canvas = hl.Canvas()

    # data = load_json()
    
model = resnet50(n_classes=num_classes)
# model.fc = torch.nn.Linear(2048, num_classes)       # 修改最后的fc层
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

criterion = nn.CrossEntropyLoss()

generate_datasets()

加载数据集

# 定义数据变换
transform = transforms.Compose([
    # transforms.Resize((64, 64)),
    transforms.ToTensor(),      # 将PILImage、nparray转换成torch
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # 归一化
])

    # 创建数据集实例
train_dataset = CustomDataset(transform=transform)

    # 创建数据加载器实例
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

test_dataset = CustomDataset(transform=transform, train=False)

    # 创建数据加载器实例
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

model = resnet50(num_classes)     # 预留的标签数量
model = model.to(DEVICE)

#原先这里选用SGD训练，但是效果很差，换成Adam优化就好了
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

valid_loader = train_dataloader
# print(list(valid_loader))    
train_acc_lst, valid_acc_lst = [], []
train_loss_lst, valid_loss_lst = [], []

训练模型

# for learn in learning_rate:
start_time = time.time()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
res = {}

for epoch in range(epochs):     # 训练10次
    
    # print(len(train_dataloader))
    num_plt = len(train_dataloader)     # 记录一共有多少训练批次
    rate = int(num_plt/50)
    
    train_loss = 0
    total_correct = 0
    train_num,test_num = 0,0
    train_loss_epoch,test_loss_epoch = 0.0,0.0      # 计算每次迭代的累计误差

    model.train()
    
    pred = []
    targ = []
    for batch_idx, (outputs, labels) in enumerate(train_dataloader):
        ### PREPARE MINIBATCH
        outputs = outputs.to(DEVICE)
        labels = labels.to(DEVICE)

        ### FORWARD AND BACK PROP
        logits, probas = model(outputs)        # 得到probas：各标签可能性
        loss = F.cross_entropy(logits, labels)     # 交叉熵函数，input[i][j]可以理解为第i ii个样本的类别为j jj的Scores
        _, predicted_labels = torch.max(probas, 1)      # 选择probas行中最大值的索引
        pred += predicted_labels.tolist()       # 这里要将每个批次的值 由tensor转变为list
        targ += labels.tolist()            # 不然会出现
        
        optimizer.zero_grad()
        loss.backward()
        
        ### 更新模型参数
        optimizer.step()
        
        train_loss += loss.item()
        train_num += outputs.size(0)
        
        # 绘制进度条
        progress_bar = '[' + ('=' * ((batch_idx + 1) // rate)) + \
                        ('*' * ((num_plt // rate - (batch_idx + 1) // rate))) + ']'# 计算当前进度等价于几个‘=’,计算‘#’后面的空格长度
                            
        print('\rschedule: {:.2f} {}  '      # \r:光标返回文本开头
                    .format(float(batch_idx/num_plt), progress_bar), end="  ")  # format：三个'{}' 分别对应可更新的 训练轮次、损失值、进度条
        
    single_loss = train_loss/len(train_dataset)
    # print(optimizer.state_dict()['param_groups'][0]['lr'])      # 从优化器中读取最终学习率 用于衰减学习率
    
    test_loss = 0
    
    cor = [0] * num_classes      # 正确的个数
    lab = [0] * num_classes      # 正确标签
    pre = [0] * num_classes      # 预测标签
    acc = [0] * num_classes      # 准确率
    rec = [0] * num_classes     # 召回率
    f1 = [0] * num_classes      # F1系数
    tot = [0] * 3

    for i, (images, labels) in enumerate(test_dataloader):
        for l in range(len(labels)):
            lab[labels[l]] += 1     # 统计实际的样本
        # print(labels.shape)
        images = images.to(device)
        labels = labels.to(device)
        logits, probas = model(images)
        loss = F.cross_entropy(logits, labels) 
        # outputs = torch.stack(probas)
        _, pred_labels = torch.max(probas, 1)      # 选择probas行中最大值得位置

        test_loss += loss.item()
        test_num += images.size(0)
        
        for j in range(len(pred_labels)):
            pre[pred_labels[j]] += 1     # 统计预测的样本
            # print(labels[j],pred_labels[j])
            if labels[j] == pred_labels[j]:
                cor[labels[j]] += 1     # 统计预测正确的样本情况
    for i in range(len(cor)):
        acc[i] = cor[i]/lab[i]      # 计算准确率
        # print(label_name[i])
        # print(acc[i])
        try:
            rec[i] = cor[i]/pre[i]      # 计算召回率
            f1[i] = 2 * acc[i] * rec[i]/ (acc[i] + rec[i]) 
        except:
            print(i)
        # print(rec[i])

        # print(f1[i])
        tot[0] += lab[i]
        tot[1] += cor[i]
        # print('accuracy of ' + label_name[i] + ': {:.5f}' .format(acc[i]))

    tot[2] = tot[1]/tot[0]

                

    model.eval()
    print('Total accuracy : {:.5f} %%'.format(accuracy_score(targ, pred)))

    if epoch not in res.keys():
        res[epoch] = {}
    res[epoch] = { "train loss": train_loss, "test loss": test_loss, "accuracy": acc}
    
    
    # 使用hiddenlayer记录日志并可视化
    train_loss_epoch = train_loss/train_num     # 训练损失函数
    test_loss_epoch = test_loss/test_num        # 测试损失函数
    history.log(epoch, train_loss = train_loss_epoch, test_loss = test_loss_epoch)
    with canvas:
        canvas.draw_plot([
            history["train_loss"],
            history["test_loss"],
        ])

    
    elapsed = (time.time() - start_time)/60
    print(f'epoch: {epoch} - learning_rate: {learning_rate} - Time elapsed: {elapsed:.2f} min')

torch.save(model, org_path) # 保存模型

分离分类错误图像、计算各类别准确率、召回率、F1参数

cor = [0] * 9
lab = [0] * 9
pre = [0] * 9
acc = [0] * 9
rec = [0] * 9
f1 = [0] * 9
tot = [0] * 3
error_folder = 'E:/1graduate_mission/ZHU/GAN/DCGAN/ResNet_pt/errorset'
for label in label_name:
    creat_path(os.path.join(error_folder,t,label))

# print(type(list(train_dataloader)[0][1]))
# for i, (images, labels) in enumerate(testset):
model = torch.load(org_path)
num = 0
for i, (images, labels) in enumerate(test_dataloader):
    # print(type(labels),labels.shape)        # <class 'torch.Tensor'> torch.Size([1])
    # print(labels.shape)
    for l in range(len(labels)):
        lab[labels[l]] += 1     # 统计实际的样本
    images = images.to(device)
    labels = labels.to(device)
    logits, probas = model(images)      #
    # loss = F.cross_entropy(logits, labels) 
    # print(logits,labels)
    # loss.backward()
    # print(labels[9])
    # outputs = torch.stack(probas)
    _, pred_labels = torch.max(probas, 1)      # 选择probas行中最大值得位置
    # pred_labels = torch.argmax(outputs, 1)
    # print(pred_labels)
    # print(labels)
    # print(images.shape)
    dim = images.shape

    for j in range(len(pred_labels)):
        pre[pred_labels[j]] += 1     # 统计预测的样本
        # print(labels[j],pred_labels[j])
        if labels[j] == pred_labels[j]:
            cor[labels[j]] += 1     # 统计预测正确的样本情况
        else:
            k = labels[j].cpu().numpy()
            # print(label_name[k])
            image = images[j,:,:,:].reshape(dim[1],dim[2],dim[3]).cpu().numpy()
            img_path = os.path.join(error_folder,t,label_name[k],f'error_{num}.png')
            cv2.imwrite(img_path,np.transpose(image*255, (1, 2, 0)))
            num += 1

for i in range(len(cor)):
    acc[i] = cor[i]/lab[i]      # 计算准确率
    try:
        rec[i] = cor[i]/pre[i]      # 计算召回率
        f1[i] = 2 * acc[i] * rec[i]/ (acc[i] + rec[i])
    except:
        print(i)
    tot[0] += lab[i]
    tot[1] += cor[i]
    # print('accuracy of ' + label_name[i] + ': {:.5f}' .format(acc[i]))

tot[2] = tot[1]/tot[0]
print(tot[2])

保存准确率于json文件中

for i in range(9):
        res[label_name[i]] = {'accuracy':acc[i],'recall':rec[i],'F1':f1[i]}

res_name = f'{t}.json'
with open(res_name, "w") as fp:  # 保存结果
    json.dump(res, fp)