本文給出了一個使用CNN＋PyTorch實現汽車電子行業(yè)視覺檢測分類詳盡的實戰(zhàn)案例解析。

在本文中，我們開發(fā)了一個卷積神經網絡（CNN），用于汽車電子行業(yè)的視覺檢測分類任務。在此過程中，我們深入研究了卷積層的概念和相關數學知識，并研究了CNN實際看到的內容以及圖像的哪些部分導致它們做出決策。

第1部分：概念背景

1.任務：將工業(yè)部件分類為合格件或廢品

在自動裝配線的一個工位中，帶有兩個突出金屬銷的線圈必須精確地定位在外殼中。金屬銷插入小插座中。在某些情況下，銷略微彎曲就會導致無法通過機器連接。視覺檢查的任務是識別這些線圈，以便可以自動將它們分類出來。

圖1：線圈、外殼和插座

為了進行檢查，每個線圈都被單獨拾起并放在屏幕前。在這個位置，相機拍攝灰度圖像。然后，由CNN檢查并分類為合格品或廢品。

圖2：視覺檢查的基本設置和生成的圖像

現在，我們要定義一個卷積神經網絡，它能夠處理圖像并從預先分類的標簽中學習。

2.什么是卷積神經網絡（CNN ）？

卷積神經網絡是卷積濾波器和全連接神經網絡（NN）的組合。CNN通常用于圖像處理，例如人臉識別或視覺檢查任務，就像我們的情況一樣。卷積濾波器是矩陣運算，它在圖像上滑動并重新計算圖像的每個像素。我們將在本文后面研究卷積濾波器。過濾器的權重不是預設的（例如Photoshop中的銳化函數），而是在訓練期間從數據中學習而來。

3.卷積神經網絡的架構

首先，讓我們來看看CNN架構的一個例子。為方便起見，我們選擇了稍后將要實現的模型。

圖3：我們的視覺檢測CNN架構

我們希望將高度為400像素、寬度為700像素的檢測圖像輸入CNN。由于圖像是灰度的，因此相應的PyTorch張量的大小為1x400x700。如果我們使用彩色圖像，我們將有3個輸入通道：一個用于紅色，一個用于綠色，一個用于藍色（RGB）。在這種情況下，張量將是3x400x700。

第一個卷積濾波器有6個大小為5x5的內核，它們在圖像上滑動并生成6個獨立的新圖像，稱為特征圖，尺寸略有縮?。?x396x696）。圖3中未明確顯示ReLU激活。它不會改變張量的維度，但會將所有負值設置為零。ReLU之后是內核大小為2x2的MaxPooling層，它將每幅圖像的寬度和高度減半。

所有三層（卷積、ReLU和MaxPooling）都是第二次實施的。這最終為我們帶來了16個特征圖，圖像高度為97像素，寬度為172像素。接下來，所有矩陣值都被展平并輸入到全連接神經網絡的大小相同的第一層中。它的第二層已經減少到120個神經元。第三層和輸出層只有2個神經元：一個代表標簽“OK”，另一個代表標簽“not OK”或“scrap”。

如果你還不清楚維度的變化，請耐心等待。我們將在下文中詳細研究不同類型的層（卷積、ReLU和MaxPooling）的工作原理及其對張量維度的影響。

4.卷積濾波器層

卷積濾波器的任務是查找圖像中的典型結構/模式。常用的內核大小為3x3或5x5。內核的9個或25個權重不是預先指定的，而是在訓練過程中學習的（這里我們假設只有一個輸入通道；否則，權重的數量將乘以輸入通道）。內核在水平和垂直方向上以定義的步幅在圖像的矩陣表示上滑動（每個輸入通道都有自己的內核）。內核和矩陣的對應值相乘并相加。每個滑動位置的求和結果形成新圖像，我們將其稱為特征圖。我們可以在卷積層中指定多個內核。在這種情況下，我們會收到多個特征圖作為結果。內核在矩陣上從左到右、從上到下滑動。因此，圖4顯示了內核在其第五個滑動位置（不計算后面的“...”部分）。我們看到紅、綠、藍（RGB）三個輸入通道。每個通道只有一個內核。在實際應用中，我們通常為每個輸入通道定義多個內核。

圖4：具有3個輸入通道和每個通道1個內核的卷積層

內核1為紅色輸入通道工作。在所示位置，我們計算特征圖中的相應新值為(-0.7)*0+(-0.9)*(-0.2)+(-0.6)*0.5+(-0.6)*0.6+0.6*(-0.3)+0.7*(-1)+0*0.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。該值放置在特征圖中以黃色突出顯示的位置。

最后，如果我們將輸入矩陣的大小與特征圖的大小進行比較，我們會發(fā)現通過核操作，我們在高度上損失了兩行，在寬度上損失了兩列。

5.ReLU激活層

卷積后，特征圖通過激活層。激活是賦予網絡非線性能力所必需的。兩種最常用的激活方法是Sigmoid和ReLU（整流線性單元）。ReLU激活將所有負值設置為零，同時保持正值不變。

圖5：特征圖的ReLU激活

在圖5中，我們看到特征圖的值逐個元素地通過了ReLU激活。

ReLU激活對特征圖的尺寸沒有影響。

6.MaxPooling層

池化層的主要任務是減小特征圖的大小，同時保留分類的重要信息。通常，我們可以通過計算內核中某個區(qū)域的平均值或返回最大值來進行池化。MaxPooling在大多數應用中更有用，因為它可以減少數據中的噪音。池化的典型內核大小為2x2或3x3。

圖6：內核為2x2的最大池化和平均池化

在圖6中，我們看到了內核大小為2x2的MaxPooling和AvgPooling的示例。特征圖被劃分為內核大小的區(qū)域，在這些區(qū)域中，我們取最大值（→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行。張量大小現在是16x194x3。

第2部分：定義和編碼CNN

從概念上講，我們已經擁有了所需的一切?，F在，讓我們進入前面1.1節(jié)中所描述的工業(yè)應用場景。

1.加載所需的庫

我們將使用幾個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")

2.配置你的設備并指定超參數

在設備中，我們存儲“cuda”或“cpu”，具體取決于你的計算機是否有可用的GPU。minibatch_size定義在模型訓練期間，一次矩陣運算將處理多少張圖像。learning_rate指定反向傳播期間參數調整的幅度，epochs定義我們在訓練階段處理整組訓練數據的頻率。

# 設備配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using {device} device")

# 指定超參數
minibatch_size = 10
learning_rate = 0.01
epochs = 60

3.自定義加載器函數

為了加載圖像，我們定義了一個custom_loader。它以二進制模式打開圖像，裁剪圖像內部的700x400像素，將其加載到內存中，并返回加載的圖像。作為圖像的路徑，我們定義相對路徑data/Coil_Vision/01_train_val_test。請確保數據存儲在你的工作目錄中。你可以從我的Dropbox下載文件??CNN_data.zip??。

#定義加載器函數
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 = "data/Coil_Vision/01_train_val_test"

4.定義數據集

我們將數據集定義為由圖像數據和標簽組成的元組，0表示廢品，1表示合格品。方法datasets.ImageFolder()從文件夾結構中讀取標簽。我們使用轉換函數首先將圖像數據加載到PyTorch張量（值介于0和1之間），然后使用近似平均值0.5和標準差0.5對數據進行歸一化。轉換后，圖像數據大致呈標準正態(tài)分布（平均值=0，標準差=1）。我們將數據集隨機分成50%的訓練數據、30%的驗證數據和20%的測試數據。

#用于加載的轉換函數
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5), (0.5))])

# 在文件夾結構中創(chuàng)建數據集
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))])

5.平衡數據集

我們的數據是不平衡的。我們的好樣本比廢棄樣本多得多。為了減少訓練期間對多數種類的偏見，我們使用WeightedRandomSampler在采樣期間為少數種類提供更高的概率。在lbls中，我們存儲訓練數據集的標簽。使用np.bincount()，我們計算0標簽（bc[0])和1標簽(bc[1]）的數量。接下來，我們計算兩個種類（p_nOK和p_OK）的概率權重，并根據lst_train列表中數據集中的順序排列它們。最后，我們從WeightedRandomSampler實例化train_sampler。

# 定義一個采樣器來平衡這些類
# 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))

6.定義數據加載器

最后，我們?yōu)橛柧?、驗證和測試數據定義三個數據加載器。數據加載器向神經網絡提供一批數據集，每批數據集由圖像數據和標簽組成。

對于train_loader和val_loader，我們將批處理大小設置為10，并對數據進行隨機打亂。test_loader使用隨機打亂數據和批處理大小1進行操作。

# 用批尺寸定義加載器
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)

7.檢查數據：繪制5個OK和5個nOK部分

為了檢查圖像數據，我們繪制了5個好樣本（“OK”）和5個廢品樣本（“nOK”）。為此，我們定義了一個2行5列的matplotlib圖形，并共享x軸和y軸。在代碼片段的核心中，我們嵌套了兩個for循環(huán)。外循環(huán)從train_loader接收數據批次。每個批次包含十張圖像和相應的標簽。內循環(huán)枚舉批次的標簽。在其主體中，我們檢查標簽是否等于0—然后我們在第二行的“nOK”下繪制圖像—或者如果標簽等于1—然后我們在第一行的“OK”下繪制圖像。一旦count_OK和count_nOK都大于或等于5，我們就中斷循環(huán)，設置標題并顯示圖形。

# 圖形和軸對象
fig, axs = plt.subplots(nrows=2, ncols=5, figsize=(20,7), sharey=True, sharex=True)

count_OK = 0
count_nOK = 0

# 循環(huán)遍歷加載器批次
for (batch_data, batch_lbls) in train_loader:

#循環(huán)遍歷batch_lbls
for i, lbl in enumerate(batch_lbls):

# 如果標簽為0 (nOK)，在第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

#如果標簽為1 (OK)，在第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

#如果兩個計數器都是>=5停止循環(huán)
if (count_OK >=5) and (count_nOK >=5):
break

# 配置繪圖畫布
fig.suptitle("Sample plot of OK and nonOK Parts", fontsize=24)
plt.setp(axs, xticks=[], yticks=[]) 
plt.show()

圖7：OK（上行）和非OK部分（下行）的示例

在圖7中，我們看到大多數nOK樣本明顯彎曲，但有些樣本肉眼無法真正區(qū)分（例如右下樣本）。

8.定義CNN模型

該模型對應于圖3中所示的架構。我們將灰度圖像（僅一個通道）輸入到第一個卷積層，并定義6個大小為5（等于5x5）的內核。卷積后跟ReLU激活和MaxPooling，內核大小為2（2x2），步長為2（2x2）。所有三個操作都以圖3中所示的尺寸重復。在__init__()方法的最后一個塊中，16個特征圖被展平并輸入到具有等效輸入大小和120個輸出節(jié)點的線性層中。它被ReLU激活，并在第二個線性層中減少到只有2個輸出節(jié)點。

在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

9.實例化模型并定義損失函數和優(yōu)化器

我們從CNN類實例化模型并將其推送到CPU或GPU上。由于我們有一個分類任務，我們選擇使用CrossEntropyLoss函數。為了管理訓練過程，我們調用隨機梯度下降（SGD）優(yōu)化器。

# 在cpu或gpu上定義模型
model = CNN().to(device)

#損失函數和優(yōu)化器
loss = nn.CrossEntropyLoss()

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

10.檢查模型的大小

為了了解模型的參數大小，我們迭代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萬個參數，其中所有參數都是可訓練的。

11.定義一個用于驗證和測試的函數

在開始模型訓練之前，讓我們準備一個函數來支持驗證和測試。函數val_test()需要一個數據加載器和CNN模型作為參數。它使用torch.no_grad()關閉梯度計算并迭代數據加載器。有了一批圖像和標簽，它將圖像輸入模型，并使用output.argmax(1)對返回的logits確定模型的預測分類。此方法返回最大值的索引；在我們的例子中，這代表分類索引。

我們計算并總結正確的預測，并保存圖像數據、預測的分類和錯誤預測的標簽。最后，我們計算準確率并將其與錯誤分類的圖像一起返回作為函數的輸出。

def val_test(dataloader, model):
# 獲取數據集大小
dataset_size = len(dataloader.dataset)

# 關閉梯度計算以進行驗證
with torch.no_grad():
# 循環(huán)數據集
correct = 0
wrong_preds = []
for (images, labels) in dataloader:
images, labels = images.to(device), labels.to(device)

#從模型中獲取原始值
output = model(images)

# 推導預測
y_pred = output.argmax(1)

# 對所有批次進行正確的分類計數
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

12.模型訓練

模型訓練由兩個嵌套的for循環(huán)組成。外循環(huán)迭代定義的epoch數，內循環(huán)枚舉train_loader。枚舉返回一批圖像數據和相應的標簽。圖像數據（images）被傳遞給模型，我們在輸出中接收模型的響應logit。outputs和真實標簽被傳遞給損失函數?；趽p失l，我們執(zhí)行反向傳播并使用optimizer.step更新參數。outputs是維度為batchsizex輸出節(jié)點的張量，在我們的例子中是10x2。我們通過行上最大值的索引（0或1）接收模型的預測。

最后，我們計算正確預測的數量（n_correct）、真正的OK部分（n_true_OK）和樣本數量（n_samples）。在每個第二個訓練周期，我們計算訓練準確率、真正的OK份額，并調用驗證函數（val_test()）。在訓練過程中，所有三個值都會被打印出來以供參考。在最后一行代碼中，我們將模型及其所有參數保存在“model.pth”中。

acc_train = {}
acc_val = {}
# 對世代進行迭代處理
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()
# 如果可用，請將數據推送到gpu
images, labels = images.to(device), labels.to(device)

#向前傳播
outputs = model(images)
l = loss(outputs, labels)

# 向后傳播和優(yōu)化
optimizer.zero_grad()
l.backward()
optimizer.step()

# 獲取預測標簽(.max返回(value,index))
_, y_pred = torch.max(outputs.data, 1)

# 計算正確的分類
n_correct += (y_pred == labels).sum().item()
n_true_OK += (labels == 1).sum().item()
n_samples += labels.size(0)

# 在世代結束時：計算準確性和打印信息
if (epoch+1) % 2 == 0:
model.eval()
# 計算準確性
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 (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}%")

# 保存模型和狀態(tài)詞典
torch.save(model, "model.pth")

在我的筆記本電腦的GPU上訓練需要幾分鐘。強烈建議從本地驅動器加載圖像；否則，訓練時間可能會增加幾個數量級！

訓練的打印輸出表明損失已顯著減少，驗證準確率（模型未用于更新其參數的數據的準確率）已達到98.4%。

如果我們繪制訓練和驗證準確率在各個時期的圖表，則可以更好地了解訓練進度。我們可以輕松做到這一點，因為我們每個第二個訓練周期都保存了值。

我們用plt.subplots()創(chuàng)建matplotlib圖和軸，并在準確率字典的鍵上繪制值。

# 實例化圖形和軸對象
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()

圖8：模型訓練期間的訓練和驗證準確率

13.加載訓練好的模型

如果你想將模型用于生產而不僅僅是用于研究目的，強烈建議你保存并加載模型及其所有參數。保存已經是訓練代碼的一部分。從驅動器加載模型同樣簡單。

#從文件中讀取模型
model = torch.load("model.pth")
model.eval()

14.使用測試數據再次檢查模型準確性

請記住，我們保留了另外20%的數據用于測試。這些數據對于模型來說是全新的，之前從未加載過。我們可以使用這些全新的數據再次檢查來驗證準確性。由于驗證數據已加載但從未用于更新模型參數，因此我們期望其準確性與測試值相似。為了進行測試，我們在test_loader上調用val_test()函數。

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

在具體示例中，我們的測試準確率達到了99.2%，但這在很大程度上取決于機會（記?。簣D像在訓練、驗證和測試數據中的隨機分布）。

15.可視化錯誤分類的圖像

錯誤分類的圖像的可視化非常簡單。首先，我們調用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)

#檢查是否發(fā)生了錯誤的預測
if len(tup[1])>=1:

# 遍歷錯誤預測的圖像
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。坦率地說，我也會將其錯誤分類。

圖9：錯誤分類的圖像

第3部分：在生產中使用訓練好的模型

1.加載模型、所需的庫和參數

在生產階段，我們假設CNN模型已經過訓練，并且參數已準備好加載。我們的目標是將新圖像加載到模型中，并讓其對相應的電子元件是否適合組裝進行分類（參見第1.1節(jié)）。

我們首先加載所需的庫，將設備設置為“cuda”或“cpu”，定義CNN類（與第2.8章完全相同），然后使用torch.load()從文件加載模型。我們需要在加載參數之前定義CNN類；否則，參數無法正確分配。

#加載所需的庫
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 = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 完全按照第2.8節(jié)來定義CNN模型
class CNN(nn.Module):

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

# 定義模型層
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

# 加載模型的參數
model = torch.load("model.pth")
model.eval()

通過運行此代碼片段，我們將CNN模型加載到計算機內存中并對其進行參數化。

2.將圖像加載到數據集中

至于訓練階段，我們需要準備圖像以供CNN模型處理。我們從指定的文件夾加載它們，裁剪內部700x400像素，并將圖像數據轉換為PyTorch張量。

#定義自定義數據集
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()的自定義數據集類中執(zhí)行所有步驟。在__init__()中，我們指定圖像文件夾，接受轉換函數，并將圖像文件夾中的圖像加載到列表self.img_lst中。方法__len__()返回圖像文件夾中的圖像數量。__getitem__()從文件夾路徑和圖像名稱組成圖像路徑，裁剪圖像的內部部分（就像我們對訓練數據集所做的那樣），并將轉換函數應用于圖像。最后，它返回圖像張量和圖像名稱。

3.路徑、轉換函數和數據加載器

數據準備的最后一步是定義一個數據加載器，允許對圖像進行迭代以進行分類。在此過程中，我們指定圖像文件夾的路徑，并將轉換函數定義為管道，首先將圖像數據加載到PyTorch張量，其次將數據規(guī)范化為大約-1到+1的范圍。我們將自定義數據集Predict_Set()實例化為變量predict_set，并定義數據加載器predict_loader。由于我們沒有指定批處理大小，predict_loader每次返回一張圖像。

# 指向圖像的路徑（最好是本地路徑，以加速加載）
path = "data/Coil_Vision/02_predict"

# 用于加載的轉換函數
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5), (0.5))])

# 創(chuàng)建數據集作為自定義數據集的實例
predict_set = Predict_Set(path, transform=transform)

# 定義加載程序
predict_loader = DataLoader(dataset=predict_set)

4.分類自定義函數

到目前為止，用于分類的圖像數據的準備工作已經完成。但是，我們仍然缺少一個自定義函數，該函數將圖像傳輸到CNN模型，將模型的響應轉換為分類，并返回分類結果。這正是我們使用predict()所做的事情。

def predict(dataloader, model):

# 關閉梯度計算
with torch.no_grad():

img_lst = []; y_pred_lst = []; name_lst = []
#循環(huán)遍歷數據加載程序
for image, name in dataloader:
img_lst.append(image)
image = image.to(device)

# 從模型中獲取原始值
output = model(image)

#推導預測
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)。圖像數據、分類結果和圖像名稱附加到列表中，列表作為函數的結果返回。

5.預測標簽和繪制圖像

最后，我們要利用自定義函數和加載器對新圖像進行分類。在文件夾“data/Coil_Vision/02_predict”中，我們保留了四張等待檢查的電子元件圖像。請記住，我們希望CNN模型告訴我們是否可以使用這些組件進行自動組裝，或者是否需要對它們進行分類，因為在嘗試將它們推入插座時，引腳可能會引起問題。

我們調用自定義函數predict()，它返回圖像列表、分類結果列表和圖像名稱列表。我們枚舉列表并以名稱和分類作為標題繪制圖像。

# 預測圖像的標簽
imgs, lbls, names  = predict(predict_loader, model)

#對分類圖像進行迭代
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()

圖10：生產階段的分類結果

我們看到左側的兩幅圖像被歸類為合格商品（標簽1），右側的兩幅圖像被歸類為廢品（標簽0）。由于我們的訓練數據，該模型非常敏感，即使針腳有輕微彎曲也會導致它們被歸類為廢品。

第4部分：CNN在“決策”中考慮了什么？

到目前為止，我們已經深入研究了CNN和我們的工業(yè)應用場景的細節(jié)。這似乎是一個很好的機會，可以更進一步，嘗試了解CNN模型在處理圖像數據時“看到”了什么。為此，我們首先研究卷積層，然后檢查圖像的哪些部分對于分類特別重要。

1.研究卷積濾波器的尺寸

為了更好地理解卷積濾波器的工作原理以及它們對圖像的作用，讓我們更詳細地檢查工業(yè)示例中的層。

要訪問這些層，我們枚舉model.children()，它是模型結構的生成器。如果該層是卷積層，我們將其附加到列表all_layers中，并將權重的維度保存在conv_weights中。如果我們有ReLU或MaxPooling層，則沒有權重。在這種情況下，我們將層和“*”附加到相應的列表中。接下來，我們枚舉all_layers，打印層類型和權重的維度。

# 設置為空的列表，以存儲圖層和權重
all_layers = []; conv_weights = []

# 迭代模型的結構
# (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("*")

# 打印層和權重維度信息
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()

圖11：層和權重的維度

請將代碼片段的輸出與圖3進行比較。第一個卷積層有一個輸入——只有一個通道的原始圖像——并返回六個特征圖。我們應用六個內核，每個內核的深度為1，大小為5x5。相應地，權重的維度為torch.Size([6,1,5,5])。相比之下，第4層接收六個特征圖作為輸入并返回16個圖作為輸出。我們應用16個卷積內核，每個內核的深度為6，大小為5x5。因此，權重的維度為torch.Size([16, 6, 5, 5])。

2.可視化卷積濾波器的權重

現在，我們知道了卷積濾波器的尺寸。接下來，我們想看看它們的權重，這是它們在訓練過程中獲得的。由于我們有如此多不同的過濾器（第一個卷積層中有6個，第二個卷積層中有16個），因此在兩種情況下，我們都選擇第一個輸入通道（索引0）。

import itertools

#遍歷所有層
for idx_out, layer in enumerate(all_layers):

#如果層是一個卷積濾波器
if type(layer) == nn.Conv2d:

# 打印層名稱
print(f"\n{idx_out+1}. Layer: {layer} \n")

# 準備繪圖并計算出權重
plt.figure(figsize=(25,6))
weights = conv_weights[idx_out][:,0,:,:] # only first input channel
weights = weights.detach().to('cpu')

# 枚舉過濾器權重（僅限第一個輸入通道）
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}")

# 打印文本
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：6+16個卷積濾波器的可視化（輸入層索引0）

圖12顯示了第1層的六個卷積濾波器內核和第4層的16個內核（用于輸入通道0）。右上角的模型示意圖用紅色輪廓表示濾波器。我們看到大多數值接近0，有些值在正或負0.20–0.25范圍內。這些數字代表圖4中卷積所使用的值。這為我們提供了接下來要檢查的特征圖。

3.檢查特征圖

根據圖4，我們通過輸入圖像的卷積獲得第一個特征圖。因此，我們從test_loader加載一個隨機圖像并將其推送到CPU（如果你在GPU上操作CNN的話）。

# 測試加載程序的批大小為1
img = next(iter(test_loader))[0].to(device)
print(f"\nImage has shape: {img.shape}\n")

# 繪制圖像
img_copy = img.to('cpu')
plt.imshow(img_copy.reshape(400,700), cmap="gray")
plt.axis("off")
plt.show()

圖13：上述代碼的輸出為隨機圖像

現在，我們將圖像數據img傳遞到第一個卷積層（all_layers[0]），并將輸出保存在results中。接下來，我們遍歷all_layers，并將上一層操作的輸出提供給下一層。這些操作是卷積、ReLU激活或MaxPoolings。我們將每個操作的輸出附加到results中。

# 將圖像通過第一層
results = [all_layers[0](img)]

# 將上一層的結果傳遞給下一層
for idx in range(1, len(all_layers)):  # Start at 1, first layer already passed!
results.append(all_layers[idx](results[-1]))  # 將最后一個結果傳遞給該圖層

最后，我們繪制原圖以及經過第一層（卷積），第二層（ReLU），第三層（MaxPooling），第四層（第二次卷積），第五層（第二次ReLU），第六層（第二次MaxPooling）之后的特征圖。

圖14：經過卷積、ReLU和MaxPooling層后的原始圖像和特征圖

我們看到卷積核（比較圖12）重新計算了圖像的每個像素。這表現為特征圖中灰度值的改變。與原始圖像相比，一些特征圖更加清晰，或者具有更強的黑白對比度，而其他特征圖似乎褪色了。

由于負值設置為零，ReLU操作將深灰色變?yōu)楹谏?/span>

MaxPooling使圖像幾乎保持不變，同時在兩個維度上將圖像大小減半。

4.可視化對分類影響最大的圖像區(qū)域

在完成之前，讓我們分析一下圖像的哪些區(qū)域對于分類為廢品（索引0）或合格部件（索引1）特別具有決定性。為此，我們使用梯度加權類激活映射（gradCAM）。該技術計算訓練模型相對于預測類別的梯度（梯度顯示輸入（圖像像素）對預測的影響程度）。每個特征圖（=卷積層的輸出通道）的梯度平均值構成了計算可視化熱圖時與特征圖相乘的權重。

但是，還是讓我們一步一步地分析一下。

def gradCAM(x):

# 運行模型并進行預測
logits = model(x)
pred = logits.max(-1)[-1] # 返回最大值（0或1）

# 在最終的conv層上獲取激活量
last_conv = model.model_layers[:5]
activations = last_conv(x)

# 計算關于模型預測的梯度
model.zero_grad()
logits[0,pred].backward(retain_graph=True)

# 計算最后一個層每個輸出通道的平均梯度
pooled_grads = model.model_layers[3].weight.grad.mean((1,2,3))

#乘以每個輸出通道與其對應的平均梯度
for i in range(activations.shape[1]):
activations[:,i,:,:] *= pooled_grads[i]

# 以所有加權輸出通道的平均值計算熱圖
heatmap = torch.mean(activations, dim=1)[0].cpu().detach()

return heatmap

我們定義一個函數gradCAM，它需要輸入數據x（圖像或特征圖），并返回熱圖。

在第一個塊中，我們在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返回的內容。

在繪制圖像和熱圖之前，我們需要對兩者進行轉換以進行疊加。請記住，特征圖比原始圖片?。▍⒁姷?.3和第1.7節(jié)），熱圖也是如此。這就是我們需要函數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將原始圖像和熱圖繪制在同一行中。

#在數據加載器上迭代
for idx, (image, name) in enumerate(predict_loader):

# 計算熱圖
image = image.to(device)
heatmap = gradCAM(image)
image = image.cpu().squeeze(0).permute(1,2,0)
heatmap = upsampleHeatmap(heatmap, image)

# 繪制圖像和熱圖
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()

圖15：圖像和熱圖（輸出的內側兩行）

熱圖中的藍色區(qū)域對模型的決策影響較小，而黃色和紅色區(qū)域則非常重要。我們發(fā)現，在我們的使用場景中，電子元件（特別是金屬針腳）的輪廓對于將零件分類為廢品或合格零件起決定性作用。當然，這是非常合理的，因為此場景中主要處理彎曲的針腳。

結論

卷積神經網絡（CNN）如今是工業(yè)環(huán)境中用于視覺檢查任務的常用且廣泛使用的工具。在我們的應用場景中，我們用相對較少的代碼行成功定義了一個模型，該模型以高精度將電子元件分類為合格零件或廢品。與傳統(tǒng)的視覺檢查方法相比，最大的優(yōu)勢是，沒有工藝工程師需要在圖像中指定視覺標記來進行分類。相反，CNN從標記的示例中學習，并能夠將這些知識復制到其他圖像中。在我們的特定場景中，626張標記圖像足以進行訓練和驗證。在更復雜的情況下，對訓練數據的需求可能會更高。

gradCAM（梯度加權類激活映射）等算法有助于理解圖像中哪些區(qū)域與模型的決策特別相關。通過這種方式，它們通過建立對模型功能的信任，支持在工業(yè)環(huán)境中廣泛使用CNN。

總之，在本文中，我們探討了卷積神經網絡內部工作原理的許多細節(jié)。希望你享受這段旅程并深入了解CNN的工作原理。

譯者介紹

朱先忠，51CTO社區(qū)編輯，51CTO專家博客、講師，濰坊一所高校計算機教師，自由編程界老兵一枚。

原文標題：??Building a Vision Inspection CNN for an Industrial Application??，作者：Ingo Nowitzky