视觉检测CNN工业应用

在本文中，我们开发并编码了一个卷积神经网络 (CNN)，用于汽车电子行业的视觉检测分类任务。在此过程中，我们深入研究了卷积层的概念和数学，并研究了 CNN 实际看到的内容以及图像的哪些部分导致它们做出决策。

PART 1：概念背景

1、任务：将工业部件分类为好部件或废品

在自动装配线的一个工位中，带有两个突出金属销的线圈必须精确地定位在外壳中。金属销插入小插座中。在某些情况下，销略微弯曲，因此无法通过机器连接。视觉检查的任务是识别这些线圈，以便可以自动将它们分类出来。

为了进行检查，每个线圈都被单独拾起并放在屏幕前。在这个位置，相机拍摄灰度图像。然后由 CNN 检查并分类为好或废品。

现在，我们要定义一个卷积神经网络，它能够处理图像并从预先分类的标签中学习。

2、什么是卷积神经网络 (CNN)？

卷积神经网络是卷积滤波器和全连接神经网络 (NN) 的组合。CNN 通常用于图像处理，例如人脸识别或视觉检查任务，就像我们的情况一样。卷积滤波器是矩阵运算，它在图像上滑动并重新计算图像的每个像素。我们将在本文后面研究卷积滤波器。过滤器的权重不是预设的（例如 Photoshop 中的锐化函数），而是在训练期间从数据中学习而来。

3、卷积神经网络的架构

让我们来看看 CNN 架构的一个例子。为方便起见，我们选择了稍后将要实现的模型。

我们希望将高度为 400 像素、宽度为 700 像素的检测图像输入 CNN。由于图像是灰度的，因此相应的 PyTorch 张量的大小为 1x400x700。如果我们使用彩色图像，我们将有 3 个输入通道：一个用于红色，一个用于绿色，一个用于蓝色（RGB）。在这种情况下，张量将是 3x400x700。

第一个卷积滤波器有 6 个大小为 5x5 的内核，它们在图像上滑动并生成 6 个独立的新图像，称为特征图，尺寸略有缩小（6x396x696）。图 3 中未明确显示 ReLU 激活。它不会改变张量的维度，但会将所有负值设置为零。ReLU 之后是内核大小为 2x2 的 MaxPooling 层。它将每幅图像的宽度和高度减半。

所有三层（卷积、ReLU 和 MaxPooling）都是第二次实施的。这最终为我们带来了 16 个特征图，图像高度为 97 像素，宽度为 172 像素。接下来，所有矩阵值都被展平并输入到全连接神经网络的大小相同的第一层中。它的第二层已经减少到 120 个神经元。第三层和输出层只有 2 个神经元：一个代表标签“OK”，另一个代表标签“不 OK”或“废品”。

如果你还不清楚维度的变化，请耐心等待。我们将在下一章中详细研究不同类型的层（卷积、ReLU 和 MaxPooling）如何工作以及它们对张量维度的影响。

4、卷积滤波器层

卷积滤波器的任务是查找图像中的典型结构/模式。常用的内核大小为 3x3 或 5x5。内核的 9 个或 25 个权重不是预先指定的，而是在训练过程中学习的（这里我们假设我们只有一个输入通道；否则，权重的数量乘以输入通道）。

内核在水平和垂直方向上以定义的步幅在图像的矩阵表示上滑动（每个输入通道都有自己的内核）。内核和矩阵的对应值相乘并相加。每个滑动位置的求和结果形成新图像，我们称之为特征图。我们可以在卷积层中指定多个内核。在这种情况下，我们收到多个特征图作为结果。内核从左到右、从上到下在矩阵上滑动。

因此，图 4 显示了内核在其第五个滑动位置（不计算“…”）。我们看到三个输入通道，分别为红色、绿色和蓝色 (RGB)。每个通道只有一个内核。在实际应用中，我们通常为每个输入通道定义多个内核。

内核 1 针对红色输入通道执行其工作。在所示位置，我们计算特征图中的相应新值为  (-0.7)0 + (-0.9)(-0.2) + (-0.6)0.5 + (-0.6)0.6 + 0.6(-0.3) + 0.7(-1) + 00.7 + (-0.1)(-0.1) + (-0.2)*(-0.1) = (-1.33)。绿色通道（内核 2）的相应计算总计为 -0.14，蓝色通道（内核 3）的相应计算总计为 0.69。为了得到特征图中特定滑动位置的最终值，我们将所有三个通道值相加并添加一个偏差（偏差和所有核权重都是在 CNN 训练期间定义的）： (-1.33) + (-0.14) + 0.69 + 0.2 = -0.58。该值放置在特征图中以黄色突出显示的位置。

最后，如果我们将输入矩阵的大小与特征图的大小进行比较，我们会发现通过核操作，我们在高度上损失了两行，在宽度上损失了两列。

5、ReLU 激活层

卷积后，特征图通过激活层。激活是赋予网络非线性能力所必需的。两种最常用的激活方法是 Sigmoid 和 ReLU（整流线性单元）。ReLU 激活将所有负值设置为零，同时保持正值不变。

在图 5 中，我们看到特征图的值逐个元素地通过了 ReLU 激活。

ReLU 激活对特征图的尺寸没有影响。

6、MaxPooling 层

池化层的主要任务是减小特征图的大小，同时保留分类的重要信息。通常，我们可以通过计算内核中某个区域的平均值或返回最大值来进行池化。MaxPooling 在大多数应用中更有用，因为它可以减少数据中的噪音。池化的典型内核大小为 2x2 或 3x3。

在图 6 中，我们看到了内核大小为 2x2 的 MaxPooling 和 AvgPooling 的示例。特征图被划分为内核大小的区域，在这些区域中，我们取最大值（→ MaxPooling）或平均值（→ AvgPooling）。

通过 2x2 核大小的池化，我们将特征图的高度和宽度减半。

7、卷积神经网络中的张量维度

现在我们已经研究了卷积滤波器、ReLU 激活和池化，我们可以修改图 3 和张量的维度。我们从 400x700 大小的图像开始。由于它是灰度的，因此只有 1 个通道，相应的张量大小为 1x400x700。我们将 6 个大小为 5x5、步幅为 1x1 的卷积滤波器应用于图像。每个滤波器都返回自己的特征图，因此我们收到其中的 6 个。由于与图 4 相比内核较大（5x5 而不是 3x3），这次我们在卷积中丢失了 4 列和 4 行。这意味着返回的张量大小为 6x396x696。

下一步，我们将具有 2x2 内核的 MaxPooling 应用于特征图（每个图都有自己的池化内核）。正如我们所了解的，这会将图的尺寸减少 2 倍。因此，张量现在的大小为 6x198x348。

现在我们应用 16 个大小为 5x5 的卷积滤波器。它们每个的内核深度为 6，这意味着每个滤波器为输入张量的 6 个通道提供单独的层。每个内核层都会在 6 个输入通道中的一个上滑动，如图 4 所示，并且 6 个返回特征图加起来为 1。到目前为止，我们只考虑了一个卷积滤波器，但我们有 16 个。这就是为什么我们收到 16 个新的特征图，每个特征图比输入小 4 列和 4 行。张量大小现在是 16x194x344。

再次，我们应用内核大小为 2x2 的 MaxPooling。由于这将特征图减半，我们现在的张量大小为 16x97x172。

最后，张量被展平，这意味着我们将所有 1697172 = 266,944 个值排列起来，并将它们输入到相应大小的全连接神经网络中。

PART 2：定义和编码 CNN
从概念上讲，我们已经拥有了所需的一切。现在，让我们进入第 1.1 章中描述的工业用例。

8、加载所需的库

我们将使用几个 PyTorch 库进行数据加载、采样和模型本身。此外，我们加载 matplotlib.pyplot 进行可视化，并加载 PIL 进行图像转换。

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.sampler import WeightedRandomSampler
from torch.utils.data import random_split
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import os
import warnings
warnings.filterwarnings("ignore")

9、配置你的设备并指定超参数

在 device中，我们存储 cuda或 cpu，具体取决于你的计算机是否有可用的 GPU。 minibatch_size定义在模型训练过程中，一次矩阵运算将处理多少张图像。 learning_rate 指定反向传播过程中参数调整的幅度， epochs 定义我们在训练阶段处理整组训练数据的频率。

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using {device} device")

# Specify hyperparameters
minibatch_size = 10
learning_rate = 0.01
epochs = 60

10、自定义加载器函数

为了加载图像，我们定义了一个 custom_loader。它以二进制模式打开图像，裁剪图像内部的 700x400 像素，将它们加载到内存中，并返回加载的图像。作为图像的路径，我们定义相对路径 data/Coil_Vision/01_train_val_test。请确保数据存储在您的工作目录中。你可以从我的 Dropbox 下载文件，文件名为 CNN_data.zip ：

# Define loader function
def custom_loader(path):
    with open(path, 'rb') as f:
        img = Image.open(f)
        img = img.crop((50, 60, 750, 460))  #Size: 700x400 px
        img.load()
        return img

# Path of images (local to accelerate loading)
path = "data/Coil_Vision/01_train_val_test"

11、定义数据集

我们将数据集定义为由图像数据和标签组成的元组，0 表示废品，1 表示好品。方法 datasets.ImageFolder() 从文件夹结构中读取标签。我们使用转换函数首先将图像数据加载到 PyTorch 张量（值介于 0 和 1 之间），然后使用近似平均值 0.5 和标准差 0.5 对数据进行归一化。转换后，图像数据大致呈标准正态分布（平均值 = 0，标准差 = 1）。我们将数据集随机分为 50% 训练数据、30% 验证数据和 20% 测试数据。

# Transform function for loading
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5), (0.5))])

# Create dataset out of folder structure
dataset = datasets.ImageFolder(path, transform=transform, loader=custom_loader)
train_set, val_set, test_set = random_split(dataset, [round(0.5*len(dataset)), 
                                                      round(0.3*len(dataset)), 
                                                      round(0.2*len(dataset))])

12、平衡数据集

我们的数据不平衡。我们的好样本比废弃样本多得多。为了减少训练期间对多数类的偏见，我们使用 WeightedRandomSampler 在采样期间为少数类提供更高的概率。在 lbls 中，我们存储训练数据集的标签。使用 np.bincount()，我们计算 0 标签 ( bc[0]) 和 1 标签 ( bc[1]) 的数量。接下来，我们计算两个类 ( p_nOK 和 p_OK) 的概率权重，并根据数据集中的顺序将它们排列在列表 lst_train 中。最后，我们从 WeightedRandomSampler 实例化 train_sampler。

# Define a sampler to balance the classes
# training dataset
lbls = [dataset[idx][1] for idx in train_set.indices]
bc = np.bincount(lbls)
p_nOK = bc.sum()/bc[0]
p_OK = bc.sum()/bc[1]
lst_train = [p_nOK if lbl==0 else p_OK for lbl in lbls]
train_sampler = WeightedRandomSampler(weights=lst_train, num_samples=len(lbls))

13、定义数据加载器

最后，我们为训练、验证和测试数据定义三个数据加载器。数据加载器向神经网络提供一批数据集，每批数据集由图像数据和标签组成。对于 train_loader 和 val_loader，我们将批大小设置为 10 并对数据进行混洗。 test_loader 使用经过打乱的数据和 1 的批处理大小进行操作。

# Define loader with batchsize
train_loader = DataLoader(dataset=train_set, batch_size=minibatch_size, sampler=train_sampler)
val_loader = DataLoader(dataset=val_set, batch_size=minibatch_size, shuffle=True)
test_loader = DataLoader(dataset=test_set, shuffle=True)

14、检查数据：绘制 5 个 OK 和 5 个 nOK 部分

为了检查图像数据，我们绘制了五个好样本（“OK”）和五个废弃样本（“nOK”）。为此，我们定义了一个有 2 行 5 列的 matplotlib 图，并共享 x 轴和 y 轴。

在代码片段的核心中，我们嵌套了两个 for 循环。外循环从 train_loader 接收批处理数据。每个批处理包含十张图像和相应的标签。内部循环枚举批次的标签。

在其主体中，我们检查标签是否等于 0 — 然后我们在第二行的“nOK”下绘制图像 — 或者如果标签等于 1 — 然后我们在第一行的“OK”下绘制图像。一旦 count_OK 和 count_nOK 都大于或等于 5，我们就中断循环，设置标题并显示图形：

# Figure and axes object
fig, axs = plt.subplots(nrows=2, ncols=5, figsize=(20,7), sharey=True, sharex=True)

count_OK = 0
count_nOK = 0

# Loop over loader batches
for (batch_data, batch_lbls) in train_loader:
    
    # Loop over batch_lbls
    for i, lbl in enumerate(batch_lbls):
        
        # If label is 0 (nOK) plot image in row 1
        if (lbl.item() == 0) and (count_nOK < 5):
            axs[1, count_nOK].imshow(batch_data[i][0], cmap='gray')
            axs[1, count_nOK].set_title(f"nOK Part#: {str(count_nOK)}", fontsize=14)
            count_nOK += 1
            
        # If label is 1 (OK) plot image in row 0
        elif (lbl.item() == 1) and (count_OK < 5):
            axs[0, count_OK].imshow(batch_data[i][0], cmap='gray')
            axs[0, count_OK].set_title(f"OK Part#: {str(count_OK)}", fontsize=14)
            count_OK += 1
    
    # If both counters are >=5 stop looping
    if (count_OK >=5) and (count_nOK >=5):
        break
        
# Config the plot canvas
fig.suptitle("Sample plot of OK and nonOK Parts", fontsize=24)
plt.setp(axs, xticks=[], yticks=[]) 
plt.show()

在图 7 中，我们看到大多数 nOK 样本明显弯曲，但有些样本无法用肉眼区分（例如右下角的样本）。

15、定义 CNN 模型

该模型对应于图 3 中所示的架构。我们将灰度图像（只有一个通道）输入第一个卷积层，并定义 6 个大小为 5（等于 5x5）的内核。卷积之后是 ReLU 激活和 MaxPooling，其内核大小为 2（2x2），步长为 2（2x2）。所有三个操作都以图 3 所示的尺寸重复。在 __init__() 方法的最后一个块中，16 个特征图被展平并输入到具有等效输入大小和 120 个输出节点的线性层中。它被 ReLU 激活，并在第二个线性层中减少到只有 2 个输出节点。

在 forward() 方法中，我们只需调用模型层并输入 x 张量。

class CNN(nn.Module):

    def __init__(self):
        super().__init__()

        # Define model layers
        self.model_layers = nn.Sequential(

            nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Flatten(),
            nn.Linear(16*97*172, 120),
            nn.ReLU(),
            nn.Linear(120, 2)
        )
        
    def forward(self, x):
        out = self.model_layers(x)
        return out

16、实例化模型并定义损失函数和优化器

我们从 CNN 类实例化 model并将其推送到 CPU 或 GPU 上。由于我们有一个分类任务，我们选择 CrossEntropyLoss 函数。为了管理训练过程，我们调用随机梯度下降 ( SGD) 优化器。

# Define model on cpu or gpu
model = CNN().to(device)

# Loss and optimizer
loss = nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

17、检查模型的大小

为了了解模型在参数方面的大小，我们迭代 model.parameters() 并首先总结所有模型参数 ( num_param)，其次总结将在反向传播期间调整的参数 ( num_param_trainable)。最后，我们打印结果。

# Count number of parameters / thereof trainable
num_param = sum([p.numel() for p in model.parameters()])
num_param_trainable = sum([p.numel() for p in model.parameters() if p.requires_grad == True])

print(f"Our model has {num_param:,} parameters. Thereof trainable are {num_param_trainable:,}!")

打印结果告诉我们该模型有超过3200万个参数，其中所有参数都是可训练的。

18、定义一个用于验证和测试的函数

在我们开始模型训练之前，让我们准备一个函数来支持验证和测试。函数 val_test()需要一个数据加载器和CNN模型作为参数。它使用 torch.no_grad()关闭梯度计算并在数据加载器上进行迭代。有了一批图像和标签，它将图像输入到模型中，并使用 output.argmax(1) 对返回的 logits 确定模型的预测类。此方法返回最大值的索引；在我们的例子中，这代表类索引。

我们计算并总结正确的预测，并保存图像数据、预测的类和错误预测的标签。最后，我们计算准确率并将其与错误分类的图像一起返回作为函数的输出。

def val_test(dataloader, model):
    # Get dataset size
    dataset_size = len(dataloader.dataset)
    
    # Turn off gradient calculation for validation
    with torch.no_grad():
        # Loop over dataset
        correct = 0
        wrong_preds = []
        for (images, labels) in dataloader:
            images, labels = images.to(device), labels.to(device)
            
            # Get raw values from model
            output = model(images)
            
            # Derive prediction
            y_pred = output.argmax(1)
            
            # Count correct classifications over all batches
            correct += (y_pred == labels).type(torch.float32).sum().item()
            
            # Save wrong predictions (image, pred_lbl, true_lbl)
            for i, _ in enumerate(labels):
                if y_pred[i] != labels[i]:
                    wrong_preds.append((images[i], y_pred[i], labels[i]))

        # Calculate accuracy
        acc = correct / dataset_size
        
    return acc, wrong_preds

19、模型训练

模型训练由两个嵌套的 for 循环组成。外循环迭代定义的 epoch 数，内循环枚举 train_loader。枚举返回一批图像数据和相应的标签。图像数据 ( images) 被传递给模型，我们在输出中接收模型的响应逻辑。输出和真实标签被传递给损失函数。根据损失 l，我们执行反向传播并使用 optimizer.step 更新参数。输出是维度为 batchsize x output nodes的张量，在我们的例子中是 10 x 2。我们通过行上最大值的索引（0 或 1）接收模型的预测。

最后，我们计算正确预测的数量 ( n_correct)、真实 OK 部分 ( n_true_OK) 和样本数量 ( n_samples)。每个第二阶段，我们计算训练准确率、真实 OK 份额，并调用验证函数 ( val_test())。在训练运行期间，所有三个值都会被打印出来以供参考。使用最后一行代码，我们将模型及其所有参数保存在 model.pth中。

acc_train = {}
acc_val = {}
# Iterate over epochs
for epoch in range(epochs):

    n_correct=0; n_samples=0; n_true_OK=0
    for idx, (images, labels) in enumerate(train_loader):
        model.train()
        # Push data to gpu if available
        images, labels = images.to(device), labels.to(device)
        
        # Forward pass
        outputs = model(images)
        l = loss(outputs, labels)
              
        # Backward and optimize
        optimizer.zero_grad()
        l.backward()
        optimizer.step()

        # Get prediced labels (.max returns (value,index))
        _, y_pred = torch.max(outputs.data, 1)

        # Count correct classifications
        n_correct += (y_pred == labels).sum().item()
        n_true_OK += (labels == 1).sum().item()
        n_samples += labels.size(0)
        
    # At end of epoch: Eval accuracy and print information
    if (epoch+1) % 2 == 0:
        model.eval()
        # Calculate accuracy
        acc_train[epoch+1] = n_correct / n_samples
        true_OK = n_true_OK / n_samples
        acc_val[epoch+1] = val_test(val_loader, model)[0]
        
        # Print info
        print (f"Epoch [{epoch+1}/{epochs}], Loss: {l.item():.4f}")
        print(f"      Training accuracy: {acc_train[epoch+1]*100:.2f}%")
        print(f"      True OK: {true_OK*100:.3f}%")
        print(f"      Validation accuracy: {acc_val[epoch+1]*100:.2f}%")
        
# Save model and state_dict
torch.save(model, "model.pth")

在我的笔记本电脑的 GPU 上训练需要几分钟。强烈建议从本地驱动器加载图像。否则，训练时间可能会增加几个数量级！

训练的打印输出表明损失已显著减少，验证准确率（模型未用于更新其参数的数据的准确率）已达到 98.4%。

如果我们绘制训练和验证准确率在各个时期的图表，则可以更好地了解训练进度。我们可以轻松做到这一点，因为我们每个第二个时期都保存了值。

我们用 plt.subplots() 创建 matplotlib 图和轴，并在准确率字典的键上绘制值：

# Instantiate figure and axe object
fig, ax = plt.subplots(figsize=(10,6))
plt.plot(list(acc_train.keys()), list(acc_train.values()), label="training accuracy")
plt.plot(list(acc_val.keys()), list(acc_val.values()), label="validation accuracy")
plt.title("Accuracies", fontsize=24)
plt.ylabel("%", fontsize=14)
plt.xlabel("Epochs", fontsize=14)
plt.setp(ax.get_xticklabels(), fontsize=14)
plt.legend(loc='best', fontsize=14)
plt.show()

20、加载训练好的模型

如果你想将模型用于生产而不仅仅是用于学习目的，强烈建议保存并加载模型及其所有参数。保存已经是训练代码的一部分。从驱动器加载模型同样简单。

# Read model from file
model = torch.load("model.pth")
model.eval()

21、使用测试数据再次检查模型准确性

请记住，我们保留了另外 20% 的数据用于测试。这些数据对于模型来说是全新的，之前从未加载过。我们可以使用这些全新的数据来再次检查验证准确性。由于验证数据已加载但从未加载过

已用于更新模型参数，我们期望准确度与测试值相似。为了进行测试，我们在 test_loader 上调用 val_test() 函数。

print(f"test accuracy: {val_test(test_loader,model)[0]*100:0.1f}%")

在具体示例中，我们的测试准确度达到 99.2%，但这在很大程度上取决于机会（记住：图像随机分布到训练、验证和测试数据）。

22、可视化错误分类的图像

错误分类的图像的可视化非常简单。首先，我们调用 val_test() 函数。它返回一个元组，其中包含索引位置 0 处的准确度值（ tup[0]），以及索引位置 1 处的另一个元组（ tup[1]），其中包含图像数据（ tup[1][0]）、预测标签（ tup[1][1]）和错误分类图像的真实标签（ tup[1][2]）。如果 tup[1] 不为空，我们将枚举它并绘制具有适当标题的错误分类图像。

%matplotlib inline

# Call test function
tup = val_test(test_loader, model)

# Check if wrong predictions occur
if len(tup[1])>=1:
    
    # Loop over wrongly predicted images
    for i, t in enumerate(tup[1]):
        plt.figure(figsize=(7,5))
        img, y_pred, y_true = t
        img = img.to("cpu").reshape(400, 700)
        plt.imshow(img, cmap="gray")
        plt.title(f"Image {i+1} - Predicted: {y_pred}, True: {y_true}", fontsize=24)
        plt.axis("off")
        plt.show()
        plt.close()
else:
    print("No wrong predictions!")

在我们的示例中，我们只有一个错误分类的图像，它代表测试数据集的 0.8%（我们有 125 张测试图像）。该图像被分类为 OK，但标签为 nOK。坦率地说，我也会将其分类错误 :)。

PART 3：在生产中使用训练好的模型

23、加载模型、所需的库和参数

在生产阶段，我们假设 CNN 模型已经过训练，并且参数已准备好加载。我们的目标是将新图像加载到模型中，并让其对相应的电子元件是否适合组装进行分类（参见第 1.1 章任务：将工业部件分类为好或废品）。

我们首先加载所需的库，将设备设置为“cuda”或“cpu”，定义 CNN 类（与第 2.8 章完全相同），然后使用 torch.load() 从文件加载模型。我们需要在加载参数之前定义 CNN 类；否则，无法正确分配参数。

# Load the required libraries
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
from PIL import Image
import os

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Define the CNN model exactly as in chapter 2.8
class CNN(nn.Module):

    def __init__(self):
        super(CNN, self).__init__()

        # Define model layers
        self.model_layers = nn.Sequential(

            nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Flatten(),
            nn.Linear(16*97*172, 120),
            nn.ReLU(),
            nn.Linear(120, 2),
            #nn.LogSoftmax(dim=1)
        )
        
    def forward(self, x):
        out = self.model_layers(x)
        return out

# Load the model's parameters
model = torch.load("model.pth")
model.eval()

通过运行此代码片段，我们将 CNN 模型加载并参数化到计算机内存中。

24、将图像加载到数据集

至于训练阶段，我们需要准备在 CNN 模型中处理的图像。我们从指定文件夹加载它们，裁剪内部 700x400 像素，并将图像数据转换为 PyTorch 张量。

# Define custom dataset
class Predict_Set(Dataset):
    def __init__(self, img_folder, transform):
        self.img_folder = img_folder
        self.transform = transform
        self.img_lst = os.listdir(self.img_folder)

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_folder, self.img_lst[idx])
        img = Image.open(img_path)
        img = img.crop((50, 60, 750, 460))  #Size: 700x400
        img.load()
        img_tensor = self.transform(img)
        return img_tensor, self.img_lst[idx]

我们执行所有在名为 Predict_Set() 的自定义数据集类中执行步骤。在 __init__()中，我们指定图像文件夹，接受转换函数，并将图像文件夹中的图像加载到列表 self.img_lst 中。方法 __len__()返回图像文件夹中的图像数量。 __getitem__()从文件夹路径和图像名称组成图像的路径，裁剪图像的内部部分（就像我们对训练数据集所做的那样），然后将 transform函数应用于图像。最后，它返回图像张量和图像名称。

25、路径、转换函数和数据加载器

数据准备的最后一步是定义一个数据加载器，允许对图像进行迭代以进行分类。在此过程中，我们指定图像文件夹的路径，并将转换函数定义为管道，首先将图像数据加载到 PyTorch 张量，然后，将数据规范化为大约 -1 到 +1 的范围。我们将自定义数据集 Predict_Set() 实例化为变量 predict_set，并定义数据加载器 predict_loader。由于我们没有指定批处理大小，因此 predict_loader 一次返回一张图像。

# Path to images (preferably local to accelerate loading)
path = "data/Coil_Vision/02_predict"

# Transform function for loading
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5), (0.5))])

# Create dataset as instance of custom dataset
predict_set = Predict_Set(path, transform=transform)

# Define loader
predict_loader = DataLoader(dataset=predict_set)

26、用于分类的自定义函数

到目前为止，用于分类的图像数据的准备工作已经完成。然而，我们仍然缺少一个自定义函数，它将图像传输到 CNN 模型，将模型的响应转换为分类，并返回分类结果。这正是我们使用 predict() 所做的事情。

def predict(dataloader, model):

    # Turn off gradient calculation
    with torch.no_grad():
        
        img_lst = []; y_pred_lst = []; name_lst = []
        # Loop over data loader
        for image, name in dataloader:
            img_lst.append(image)
            image = image.to(device)
            
            # Get raw values from model
            output = model(image)
            
            # Derive prediction
            y_pred = output.argmax(1)
            y_pred_lst.append(y_pred.item())
            name_lst.append(name[0])
            
    return img_lst, y_pred_lst, name_lst

predict() 需要数据加载器和 CNN 模型作为其参数。在其核心中，它迭代数据加载器，将图像数据传输到模型，并使用 output.argmax(1) 将模型响应解释为分类结果 — 0 表示废品部件 ( nOK)，1 表示合格部件 ( OK)。图像数据、分类结果和图像名称被附加到列表中，列表作为函数的结果返回。

27、预测标签和绘制图像

最后，我们要利用自定义函数和加载器对新图像进行分类。在文件夹 data/Coil_Vision/02_predict中，我们保留了四张等待检查的电子元件图像。请记住，我们希望 CNN 模型告诉我们是否可以使用这些组件进行自动组装，或者是否需要对它们进行分类，因为在尝试将它们推入插座时，引脚可能会引起问题。
我们调用自定义函数 predict()，它返回一个图像列表、一个分类结果列表和一个图像名称列表。我们枚举列表并绘制图像，并以名称和分类作为标题：

# Predict labels for images
imgs, lbls, names  = predict(predict_loader, model)

# Iterate over classified images
for idx, image in enumerate(imgs):
    plt.figure(figsize=(8,6))
    plt.imshow(image.squeeze(), cmap="gray")
    plt.title(f"\nFile: {names[idx]}, Predicted label: {lbls[idx]}", fontsize=18)
    plt.axis("off")
    plt.show()
    plt.close()

我们看到左侧的两幅图像被归类为商品（标签 1），右侧的两幅图像被归类为废品（标签 0）。由于我们的训练数据，该模型非常敏感，即使针脚有轻微弯曲也会导致它们被归类为废品。

PART 4：CNN 在“决策”中考虑了什么？

到目前为止，我们已经深入研究了 CNN 和我们的工业用例的细节。这似乎是一个很好的机会，可以更进一步，尝试了解 CNN 模型在处理图像数据时“看到”了什么。为此，我们首先研究卷积层，然后检查图像的哪些部分对于分类特别重要。

28、研究卷积滤波器的尺寸

为了更好地理解卷积滤波器的工作原理以及它们对图像的作用，让我们更详细地检查工业示例中的层。

要访问这些层，我们枚举 model.children()，它是模型结构的生成器。如果该层是卷积层，我们将其附加到列表 all_layers 中，并将权重的维度保存在 conv_weights 中。如果我们有 ReLU 或 MaxPooling 层，则没有权重。在这种情况下，我们将层和 *​​附加到相应的列表中。接下来，我们枚举 all_layers，打印层类型和权重的维度。

# Empty lists to store the layers and the weights
all_layers = []; conv_weights = []

# Iterate over the model's structure
# (First level nn.Sequential)
for _, layer in enumerate(list(model.children())[0]):
    if type(layer) == nn.Conv2d:
        all_layers.append(layer)
        conv_weights.append(layer.weight)
    elif type(layer) in [nn.ReLU, nn.MaxPool2d]:
        all_layers.append(layer)
        conv_weights.append("*")

# Print layers and dimensions of weights
for idx, layer in enumerate(all_layers):
    print(f"{idx+1}. Layer: {layer}")
    if type(layer) == nn.Conv2d:
        print(f"          weights: {conv_weights[idx].shape}")
    else:
        print(f"          weights: {conv_weights[idx]}")
    print()

请将代码片段的输出与图 3 进行比较。第一个卷积层有一个输入 — 只有一个通道的原始图像 — 并返回六个特征图。我们应用六个内核，每个内核的深度为 1，大小为 5x5。相应地，权重的维度为 torch.Size([6, 1, 5, 5])。相比之下，第 4 层接收六个特征图作为输入并返回 16 个图作为输出。我们应用 16 个卷积内核，每个内核的深度为 6，大小为 5x5。因此，权重的维度为 torch.Size([16, 6, 5, 5])。

29、可视化卷积滤波器的权重

现在，我们知道了卷积滤波器的尺寸。接下来，我们想看看它们的权重，这是它们在训练过程中获得的。由于我们有许多不同的过滤器（第一个卷积层中有六个，第二个卷积层中有 16 个），因此在两种情况下，我们都选择第一个输入通道（索引 0）。

import itertools

# Iterate through all layers
for idx_out, layer in enumerate(all_layers):
    
    # If layer is a convolutional filter
    if type(layer) == nn.Conv2d:
        
        # Print layer name
        print(f"\n{idx_out+1}. Layer: {layer} \n")
        
        # Prepare plot and weights
        plt.figure(figsize=(25,6))
        weights = conv_weights[idx_out][:,0,:,:] # only first input channel
        weights = weights.detach().to('cpu')
        
        # Enumerate over filter weights (only first input channel)
        for idx_in, f in enumerate(weights):
            plt.subplot(2,8, idx_in+1)
            plt.imshow(f, cmap="gray")
            plt.title(f"Filter {idx_in+1}")
            
            # Print texts
            for i, j in itertools.product(range(f.shape[0]), range(f.shape[1])):
                if f[i,j] > f.mean():
                    color = 'black'
                else:
                    color = 'white'
                plt.text(j, i, format(f[i, j], '.2f'), horizontalalignment='center', verticalalignment='center', color=color)
            
            plt.axis("off")
        plt.show()
        plt.close()       

我们遍历 all_layers。如果该层是卷积层 ( nn.Conv2d)，则我们打印该层的索引和该层的核心数据。接下来，我们准备一个图并提取第一个输入层的权重矩阵作为示例。我们枚举所有输出层并使用 plt.imshow() 绘制它们。最后，我们在图像上打印权重的值，以便我们直观地可视化卷积滤波器。

图 12 显示了第 1 层的六个卷积滤波器内核和第 4 层的 16 个内核（用于输入通道 0）。右上角的模型示意图用红色轮廓表示滤波器。我们看到大多数值接近 0，有些值在正或负 0.20–0.25 范围内。这些数字表示图 4 中演示的卷积所使用的值。这为我们提供了特征图，我们接下来将对其进行检查。

30、检查特征图

根据图 4，我们通过输入图像的卷积获得了第一个特征图。因此，我们从 test_loader 加载一个随机图像并将其推送到 CPU（如果你在 GPU 上操作 CNN）：

# Test loader has a batch size of 1
img = next(iter(test_loader))[0].to(device)
print(f"\nImage has shape: {img.shape}\n")

# Plot image
img_copy = img.to('cpu')
plt.imshow(img_copy.reshape(400,700), cmap="gray")
plt.axis("off")
plt.show()

现在我们将图像数据 img 传递到第一个卷积层 ( all_layers[0])，并将输出保存在 results 中。接下来，我们遍历 all_layers，并将上一层操作的输出提供给下一层。这些操作是卷积、ReLU 激活或 MaxPoolings。我们将每个操作的输出附加到 results 中：

# Pass the image through the first layer
results = [all_layers[0](img)]

# Pass the results of the previous layer to the next layer
for idx in range(1, len(all_layers)):  # Start at 1, first layer already passed!
    results.append(all_layers[idx](results[-1]))  # Pass the last result to the layer

最后，我们绘制原始图像，即通过第一层后的特征图，第一层（卷积）、第二层（ReLU）、第三层（MaxPooling）、第四层（第二次卷积）、第五层（第二次 ReLU）和第六层（第二次 MaxPooling）。

我们看到卷积核（比较图 12）重新计算了图像的每个像素。这表现为特征图中灰度值的变化。与原始图像相比，一些特征图更加锐化或具有更强的黑白对比度，而其他特征图似乎褪色了。

由于负值设置为零， ReLU 操作将深灰色变成黑色。

MaxPooling 使图像几乎保持不变，同时在两个维度上将图像大小减半。

31、可视化对分类影响最大的图像区域

在完成之前，让我们分析一下图像的哪些区域对于分类为废品（索引 0）或好部件（索引 1）特别具有决定性。为此，我们使用梯度加权类激活映射 ( gradCAM)。该技术计算训练模型相对于预测类的梯度（梯度显示输入（图像像素）对预测的影响程度）。每个特征图（= 卷积层的输出通道）的梯度平均值构成了在计算可视化热图时与特征图相乘的权重。

但让我们一步一步地看：

def gradCAM(x):
    
    # Run model and predict
    logits = model(x)
    pred = logits.max(-1)[-1] # Returns index of max value (0 or 1)
    
    # Fetch activations at final conv layer
    last_conv = model.model_layers[:5]
    activations = last_conv(x)

    # Compute gradients with respect to model's prediction
    model.zero_grad()
    logits[0,pred].backward(retain_graph=True)
    
    # Compute average gradient per output channel of last conv layer
    pooled_grads = model.model_layers[3].weight.grad.mean((1,2,3))
    
    # Multiply each output channel with its corresponding average gradient
    for i in range(activations.shape[1]):
        activations[:,i,:,:] *= pooled_grads[i]
    
    # Compute heatmap as average over all weighted output channels
    heatmap = torch.mean(activations, dim=1)[0].cpu().detach()
    
    return heatmap

我们定义一个函数 gradCAM，它需要输入数据 x、图像或特征图，并返回 heatmap热图。

在第一个块中，我们在 CNN 模型中输入 x 并接收 logits，这是一个形状为 [1, 2] 的张量，只有两个值。这些值表示类 0 和 1 的预测概率。我们选择较大值的索引作为模型的预测 pred。

在第二个块中，我们提取模型的前五层 - 从第一个卷积到第二个 ReLU - 并将它们保存到 last_conv。我们在选定的层中运行 x 并将输出存储在激活中。顾名思义，这些是第二个卷积层（ReLU 激活后）的激活（=特征图）。

在第三个块中，我们对预测类 logits[0,pred] 的 logit 值进行反向传播。换句话说，我们计算 CNN 相对于预测的所有梯度。梯度显示输入数据（原始图像像素）的变化对模型输出（预测）的影响程度。结果保存在 PyTorch 计算图中，直到我们使用 model.zero_grad() 将其删除。

在第四个块中，我们计算输入通道上的梯度平均值，以及图像或特征图的高度和宽度。结果，我们收到从第二个卷积层返回的 16 个特征图的 16 个平均梯度。我们将它们保存在 pooled_grads 中。

在第五个块中，我们迭代从第二个卷积层返回的 16 个特征图，并使用平均梯度 pooled_grads 对它们进行加权。此操作对那些对预测具有高重要性的特征图（及其像素）产生更大的影响，反之亦然。从现在开始，激活不再保留特征图，而是保留加权特征图。

最后，在最后一个块中，我们将热图计算为所有激活的平均特征图。这就是函数 gradCAM 返回的内容。

在绘制图像和热图之前，我们需要对两者进行转换以进行叠加。请记住，特征图小于原始图片（参见第 1.3 章和第 1.7 章），热图也是如此。这就是我们需要函数 upsampleHeatmap() 的原因。该函数将像素值缩放到 0 到 255 的范围，并将它们转换为 8 位整数格式（cv2 库需要）。它将热图的大小调整为 400x700 像素，并将颜色图应用于图像和热图。最后，我们叠加 70% 的热图和 30% 的图像并返回绘图的合成图。

import cv2

def upsampleHeatmap(map, img):
    m,M = map.min(), map.max()
    i,I = img.min(), img.max()
    map = 255 * ((map-m) / (M-m))
    img = 255 * ((img-i) / (I-i))
    map = np.uint8(map)
    img = np.uint8(img)
    map = cv2.resize(map, (700,400))
    map = cv2.applyColorMap(255-map, cv2.COLORMAP_JET)
    map = np.uint8(map)
    img = cv2.applyColorMap(255-img, cv2.COLORMAP_JET)
    img = np.uint8(img)
    map = np.uint8(map*0.7 + img*0.3)
    return map

我们希望将原始图像和热图叠加层并排绘制在同一行中。为此，我们迭代数据加载器 predict_loader，在图像上运行 gradCAM() 函数，在热图和图像上运行 upsampleHeatmap() 函数。最后，我们使用 matplotlib.pyplot 将原始图像和热图绘制在同一行中。

# Iterate over dataloader
for idx, (image, name) in enumerate(predict_loader):
    
    # Compute heatmap
    image = image.to(device)
    heatmap = gradCAM(image)
    image = image.cpu().squeeze(0).permute(1,2,0)
    heatmap = upsampleHeatmap(heatmap, image)
    
    # Plot images and heatmaps
    fig = plt.figure(figsize=(14,5))
    fig.suptitle(f"\nFile: {names[idx]}, Predicted label: {lbls[idx]}\n", fontsize=24)
    plt.subplot(1, 2, 1)
    plt.imshow(image, cmap="gray")
    plt.title(f"Image", fontsize=14)
    plt.axis("off")
    plt.subplot(1, 2, 2)
    plt.imshow(heatmap)
    plt.title(f"Heatmap", fontsize=14)
    plt.tight_layout()
    plt.axis("off")
    plt.show()
    plt.close()

热图的蓝色区域对模型的决策影响较小，而黄色和红色区域则非常重要。我们看到，在我们的用例中，电子元件（特别是金属针脚）的轮廓对于分类为废品或好部件起着决定性作用。当然，这是非常合理的，因为用例主要处理弯曲的针脚。

32、结束语

卷积神经网络 (CNN) 如今是工业环境中用于视觉检查任务的常用且广泛使用的工具。在我们的用例中，我们用相对较少的代码行成功定义了一个模型，该模型以高精度将电子元件分类为好部件或废品。与传统的视觉检测方法相比，最大的优势在于，工艺工程师无需在图像中指定视觉标记来进行分类。相反，CNN 从标记的示例中学习，并能够将这些知识复制到其他图像中。在我们的特定用例中，626 张标记图像足以进行训练和验证。在更复杂的情况下，对训练数据的需求可能会高得多。

像 gradCAM（梯度加权类激活映射）这样的算法有助于理解图像中哪些区域与模型的决策特别相关。通过这种方式，它们通过建立对模型功能的信任，支持在工业环境中广泛使用 CNN。

在本文中，我们探讨了卷积神经网络内部工作原理的许多细节。我希望你喜欢这段旅程，并深入了解 CNN 的工作原理。

原文链接：Building a Vision Inspection CNN for an Industrial Application

汇智网翻译整理，转载请标明出处