基于塊金字塔記憶模塊的無監(jiān)督異常檢測

鄢 寧，李岳陽，羅海馳

（1.江南大學(xué) 人工智能與計算機(jī)學(xué)院，江蘇 無錫 214122；2.江南大學(xué) 物聯(lián)網(wǎng)工程學(xué)院，江蘇 無錫 214122）

0 概述

異常檢測主要是判斷輸入數(shù)據(jù)是否與訓(xùn)練數(shù)據(jù)“類似”，主要是針對少數(shù)、不可預(yù)測或不確定的罕見的事件。在現(xiàn)實(shí)世界中，由于異常種類繁多，出現(xiàn)的概率較低，且形式多種多樣，因此難以收集到大量的、所有種類的異常樣本。但是人們可以很容易地獲取到大量的正常樣本，這將導(dǎo)致正負(fù)樣本的比例極其不均衡，另外手工標(biāo)注工作也非常耗時。因此，無監(jiān)督異常檢測［1-3］的方法至關(guān)重要，在訓(xùn)練時，通過使用正常樣本進(jìn)行訓(xùn)練從而得到正常樣本的數(shù)據(jù)分布；在推斷時，具有與正常樣本不同分布的數(shù)據(jù)樣本被認(rèn)為是異常樣本。

在無監(jiān)督異常檢測方法中，基于重建的方法在訓(xùn)練過程中只需要正常樣本，不需要額外的訓(xùn)練數(shù)據(jù)和預(yù)訓(xùn)練模型［4］，因此被廣泛應(yīng)用。如基于自編碼器［5-7］或生成對抗網(wǎng)絡(luò)［8-10］的重建方法已經(jīng)得到了廣泛的研究，這種方法是基于正常樣本比異常樣本能夠通過潛在特征空間更好地被重建這一假設(shè)［11-12］，使重建網(wǎng)絡(luò)在正常圖像上進(jìn)行訓(xùn)練后能夠充分地學(xué)習(xí)到正常圖像的特征，從而在測試時能夠很好地重建正常圖像。但是，對于異常圖像來說，由于在訓(xùn)練階段沒有學(xué)習(xí)異常特征，因此在測試時其重建能力較差，需通過計算輸入圖像和重建圖像之間的誤差來進(jìn)行異常檢測。但是這一假設(shè)并不總是成立，由于卷積神經(jīng)網(wǎng)絡(luò)具有強(qiáng)大的泛化能力，有時異常特征也能夠被重建出來，因此無法通過比較異常圖像和其重建圖像的差異性來判斷是否存在異常。

為了解決上述問題，文獻(xiàn)［13-15］通過在自編碼器中增加記憶模塊，并在訓(xùn)練階段用于存儲正常樣本的典型模式，在測試階段通過聚合內(nèi)存中與輸入相關(guān)的正常項(xiàng)來減少重建異常。但是，這種方式容易忽略圖像細(xì)節(jié)，導(dǎo)致輸出圖像模糊，且原始的記憶模塊［13］是以逐像素的方式進(jìn)行存儲和讀取的，由于像素點(diǎn)的模式變化是有限的，異常像素點(diǎn)和正常像素點(diǎn)可能會具有相同的模式，那么異常像素點(diǎn)就很容易被重建，最終導(dǎo)致重建異常，因此該方法性能提升有限。

文獻(xiàn)［16］提出記憶模塊使用合適的塊進(jìn)行存儲，異常塊與正常塊很可能會具有不同塊模式，從而對異常塊進(jìn)行較差的重建。為了得到更好的檢測結(jié)果，選擇合適的塊的尺寸就顯得尤為重要。如果塊的尺寸過大，那么模型重建的圖像會丟失很多細(xì)節(jié)信息，從而導(dǎo)致重建模糊；如果塊的尺寸過小，異常就會被重建，從而導(dǎo)致檢測精度降低。此外，對于不同種類的數(shù)據(jù)集，由于其特點(diǎn)不同，合適的塊的尺寸也不同。對于某個種類的數(shù)據(jù)集，其需要檢測的異常的面積大小差距較大，對于面積較大的異常，如果選用較小尺寸的塊存儲和讀取，則容易重建異常；對于面積較小的異常，如果選用較大尺寸的塊存儲和讀取，則容易重建模糊。為此，需要進(jìn)行大量的實(shí)驗(yàn)以得到合適塊的尺寸。

基于上述問題，本文提出一種基于改進(jìn)的塊記憶模塊的無監(jiān)督異常檢測方法。在原有的塊記憶模塊［16］中加入多個尺度的塊記憶模塊，各尺度塊記憶模塊通過讀取、聚合得到多特征圖融合的輸出特征圖，最大限度地保留不同尺度下樣本的特征信息，以更好地重建數(shù)據(jù)，減少重建異常的情況。通過增加skip connection［17］結(jié)構(gòu)，將編碼器輸出的多個尺度的特征圖使用塊金字塔記憶模塊得到多個增強(qiáng)特征圖，將skip connection 送入解碼器進(jìn)行重建，避免直接將異常特征送入解碼器進(jìn)行重建，以在提高重建圖像的清晰度的同時減少重建異常的情況。最終采用SSIM［18］損失，通過亮度、對比度和結(jié)構(gòu)3 個維度來增強(qiáng)圖像重建效果，并且作為異常分?jǐn)?shù)組成部分，增強(qiáng)原始圖像和重建圖像的差異性。

1 相關(guān)工作

1.1 基于重建的無監(jiān)督異常檢測

基于重建的方法通常使用自編碼器［19-21］進(jìn)行圖像重建。自編碼器由編碼器和解碼器組成，編碼器將輸入圖像映射到低維特征空間，解碼器試圖從低維特征空間中恢復(fù)數(shù)據(jù)，獲得重建圖像。文獻(xiàn)［8，22］引入對抗學(xué)習(xí)判別器，使用對抗訓(xùn)練來提高圖像重建質(zhì)量。除了使用基于圖像的重建損失和對抗損失來訓(xùn)練自編碼器外，最小化原始圖像和重建圖像在潛在特征空間之間的距離也有助于學(xué)習(xí)正常樣本的數(shù)據(jù)分布。在推理階段，除計算原始圖像和重建圖像的差異分?jǐn)?shù)外，原始圖像和重建圖像潛在特征空間距離的分?jǐn)?shù)在異常評判中也發(fā)揮重要的作用。

1.2 記憶模塊

由于卷積神經(jīng)網(wǎng)絡(luò)強(qiáng)大的泛化能力，且有些異常與正常特征具有相同的局部模式，異常也能夠被重建。為此，文獻(xiàn)［13］提出在自編碼器中增加記憶模塊，將編碼器的輸出不直接傳入解碼器，而是作為查詢來檢索內(nèi)存中相關(guān)的項(xiàng)，然后將這些項(xiàng)聚合后送入解碼器。由于記憶模塊存儲的是正常數(shù)據(jù)的典型模式，因此在測試階段，從內(nèi)存中獲取正常的數(shù)據(jù)信息用于重建，來增強(qiáng)異常的重建誤差。

記憶模塊如圖1 所示，其中內(nèi)存M?RN×C包含N個維度為C的實(shí)值向量，mi表示內(nèi)存M的第i個內(nèi)存項(xiàng)，mi?RC,?i?{1,2,…,N}。內(nèi)存M是 可尋址內(nèi)存，尋址方案是根據(jù)查詢與內(nèi)存項(xiàng)的相似性來計算注意權(quán)值w，w?R1×C。查詢z關(guān)于內(nèi)存M中的內(nèi)存項(xiàng)mi的權(quán)值wi的計算公式如下：

圖1 記憶模塊網(wǎng)絡(luò)結(jié)構(gòu)Fig.1 Memory module network structure

然后通過使用尋址權(quán)值w對內(nèi)存M中的內(nèi)存項(xiàng)進(jìn)行線性組合得到?RC，計算公式如下：

但是一些異常仍然能夠被重建，為了減少這種情況，并且避免受不重要信息的影響，對w進(jìn)行稀疏化操作，計算公式如下：

其中：λ為收縮閾值（在實(shí)驗(yàn)中將其設(shè)置為1/N）；?是一個非常小的正參量；max(?,0)為ReLU 激活函數(shù)。然后進(jìn)行歸一化操作，得到通過稀疏化操作，使用少量且更相關(guān)的內(nèi)存項(xiàng)來聚合得 到。

2 本文方法

2.1 整體網(wǎng)絡(luò)結(jié)構(gòu)

提出網(wǎng)絡(luò)結(jié)構(gòu)由重建子網(wǎng)絡(luò)G和判別子網(wǎng)絡(luò)D組 成，網(wǎng)絡(luò)結(jié) 構(gòu)如圖2 所 示。采 用UNet［17］作為基 礎(chǔ)重建子網(wǎng)絡(luò)，其由編碼器GE和解碼器GD組成。給定輸入圖像x?Rh×w×c，其中，h為高，w為寬，c為通道維數(shù)。編碼器GE將x映射到低維特征空間，解碼器GD試圖從投影的低維空間中恢復(fù)數(shù)據(jù)得到重建圖像?Rh×w×c，其中，h為高，w為寬，c為通道維數(shù)。為了減少圖像重建模糊和重建異常的情況，在編碼器GE和解碼器GD之間增加skip connection 和塊金字塔記憶模塊得到改進(jìn)后的重建子網(wǎng)絡(luò)G，將編碼器輸出的特征圖作為查詢檢索內(nèi)存中相關(guān)的內(nèi)存項(xiàng)，聚合后通過skip connection 傳入解碼器用于重建。

圖2 整體網(wǎng)絡(luò)結(jié)構(gòu)Fig.2 Overall network structure

判別網(wǎng)絡(luò)D的網(wǎng)絡(luò)結(jié)構(gòu)是按照DCGAN［23］的判別網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計的，引入判別網(wǎng)絡(luò)的目的主要有：1）使用對抗學(xué)習(xí)思想提升重建精度，在訓(xùn)練階段，判別網(wǎng)絡(luò)D用來預(yù)測給定輸入的類標(biāo)簽，對輸入圖像x和重建圖像進(jìn)行分類，將輸入圖像x判別為真，重建圖像判別為假；2）判別網(wǎng)絡(luò)D除作為分類器外，還作為特征提取器，提取輸入圖像x和重建圖像的潛在特征表示，再計算兩者之間的差異性，訓(xùn)練階段通過最小化這一差異性來進(jìn)一步提升重建效果；在測試階段，將其作為異常評價指標(biāo)的組成部分提升檢測精度。

2.2 塊金字塔記憶模塊

對于輸入圖像x，編碼器GE輸出一組特征，記為{F1,F2,…,FL}，其中，L表示編碼器GE共輸出L層特征圖。本文將每層特征圖作為其對應(yīng)的塊金字塔記憶模塊的輸入，對其進(jìn)行不同尺度的存儲和讀取，以減少異常重建。

對于編碼器GE輸出的第l層特征圖其中，Hl為高，Wl為寬，Cl為通道維數(shù)。對其使用塊金字塔記憶模塊進(jìn)行存儲和讀取如圖3所示。

圖3 塊金字塔記憶模塊網(wǎng)絡(luò)結(jié)構(gòu)Fig.3 Block pyramid memory module network structure

對于第i個分支，將其均分為個塊，其 中，rHi、rWi和rCi分別為 沿高度、寬度和 通道維度的劃分率。每個塊的高度為hli=rHi×Hl，寬度為wli=rWi×Wl，通道長度為cli=rCi×Cl。對于每一個塊，將其展平后得到一個向量作為查詢，檢索內(nèi)存（表示內(nèi)存Mli包含Nli個維度為Pli的實(shí)值向量，Pli=hli×wli×cli）中相關(guān)的內(nèi)存項(xiàng)，再經(jīng)過聚合得到的特征向量，最后變換為與Zl相同的維度大小的特征圖

本文通過增加多個尺度的塊記憶模塊來減少重建模糊和重建異常的情況。分支的劃分率設(shè)置如下：

2.3 損失函數(shù)

本文方法只使用正常圖像進(jìn)行訓(xùn)練，在測試階段同時使用正常圖像和異常圖像。因?yàn)楫惓D像沒有參與訓(xùn)練，所以在測試時異常圖像與其重建圖像在圖像空間和潛在的特征空間之間的差異性較大。本文結(jié)合3 個損失函數(shù)來訓(xùn)練模型，以保證模型可以更好地重建正常圖像。

1）第1 項(xiàng)為重建損失，其中包括L1 損失和基于塊的SSIM［18］損失。

L1 損失確保模型能夠?qū)斎雸D像x進(jìn)行精確的重建，通過把輸入圖像x與重建圖像（即G(x)）的絕對差值的總和最小化來優(yōu)化G：

基于塊的SSIM 損失的計算公式如下：

其中：h和w為圖像x（或者）的高和寬；NP為x（或者）的像素總個數(shù)；SSIM(x,)i,j表示輸入圖像x與重建圖像中以(i,j)為中心的兩個塊的SSIM 值。

因此，重建損失為：

其中：λ為SSIM 損失的權(quán)重參數(shù)。

2）第2 項(xiàng)損失為潛在特征空間損失，為確保輸入圖像x和重建圖像有盡可能相似的潛在特征表示，使用判別網(wǎng)絡(luò)D最后一個卷積層的輸出來提取輸入圖像x和重建圖像的潛在特征表示，令z=f(x)和=f()，采用L2 損失：

3）第3 項(xiàng)為對抗損失［24］，使用判別網(wǎng)絡(luò)進(jìn)行對抗訓(xùn)練，采用對抗損失來使生成網(wǎng)絡(luò)G盡可能逼真地重建正常圖像，判別網(wǎng)絡(luò)D對輸入圖像x與重建圖像進(jìn)行分類，采用二分類的交叉熵?fù)p失來進(jìn)行優(yōu)化：

因此，總的訓(xùn)練損失為上述3 個損失的加權(quán)總和：

其中：α、β和γ是調(diào)節(jié)單個損失對總體目標(biāo)函數(shù)影響的權(quán)重參數(shù)。

2.4 異常分?jǐn)?shù)

對于測試圖像I，其異常分?jǐn)?shù)A(I)為：

其中：a為權(quán)重參數(shù)。

根據(jù)式（10）計算出測試集中每個測試樣本的異常分?jǐn)?shù)，將測試集中所有樣本的異常分?jǐn)?shù)表示為向量A，并采用文獻(xiàn)［9］方式將其縮放到[0,1]內(nèi)，因此測試樣本I的最終異常得分為：

3 實(shí)驗(yàn)

3.1 數(shù)據(jù)集和實(shí)驗(yàn)設(shè)置

本文使用工業(yè)異常檢測數(shù)據(jù)集MVTec AD［25］，該數(shù)據(jù)集專為無監(jiān)督工業(yè)視覺異常檢測而開發(fā)，是目前該領(lǐng)域的重要基準(zhǔn)，包括用來訓(xùn)練的3 629 張正常圖像和用于測試的1 725 張圖像（包括正常圖像和異常圖像），由15 種類別的對象組成，包括10 種物體類和5 種紋理類，每類圖像包括一個訓(xùn)練集和一個測試集。圖像的分辨率在700～1 024 像素范圍內(nèi)，該數(shù)據(jù)集的挑戰(zhàn)在于有些異常圖像與正常圖像的差別較小，且異常面積的大小差別較大。在實(shí)驗(yàn)時，本文將圖像大小調(diào)整為256×256 像素。

本文方法是使用PyTorch 1.8.1 和CUDA 11.1 實(shí)現(xiàn)的，所有實(shí)驗(yàn)均在RTX 3090 GPU 上運(yùn)行。使用Adam［26］優(yōu)化器，設(shè) 置lr=0.000 1，β1=0.500，β2=0.999，weight-decay=0，Epoch 為800，batch_size=8。所有對比實(shí)驗(yàn)和消融實(shí)驗(yàn)均在同一環(huán)境下完成。

3.2 評價指標(biāo)

本文實(shí)驗(yàn)采用AUC 來評價模型的性能，當(dāng)AUC分?jǐn)?shù)越高時，模型性能越好。AUC 表示隨機(jī)抽取一對樣本（一個正樣本和一個負(fù)樣本），然后用訓(xùn)練得到的分類器來對這兩個樣本進(jìn)行預(yù)測，該正樣本預(yù)測值大于該負(fù)樣本預(yù)測值的概率為：

其中：M表示正樣本個數(shù)；N表示負(fù)樣本個數(shù)；rank(i)表示第i個樣本的序號（M個正樣本按照預(yù)測值從小到大進(jìn)行排序）。

3.3 主流方法和本文方法對比

為驗(yàn)證本文方法的有效性，將主流基于重建的無監(jiān)督異常檢測方法GANomaly［8］、Skip-GANomaly［9］、MemAE［13］和DAAD+［16］與本文方法在MVTec AD 數(shù)據(jù)集上進(jìn)行對比，實(shí)驗(yàn)結(jié)果如表1 所示，其中加粗?jǐn)?shù)字為最優(yōu)值。

表1 主流方法模型與本文改進(jìn)方法的AUC 對比Table 1 Comparison of AUC between mainstream method models and improved method in this paper

分析表1 的數(shù)據(jù)可知，本文方法明顯優(yōu)于主流的基于重建的無監(jiān)督異常檢測方法。GANomaly 存在的重建模糊問題和Skip-GANomaly 存在的重建異常問題導(dǎo)致了兩者性能均不佳。MemAE 方法以逐像素點(diǎn)的方式存儲正常樣本的典型模式，在進(jìn)行異常檢測時，輸入特征圖的每個像素點(diǎn)都通過加權(quán)聚合相似的內(nèi)存項(xiàng)得到。由于異常像素點(diǎn)與正常像素點(diǎn)可能具有相同的模式，異常很容易被重建，因此性能提升有限。而本文方法是以記憶塊的方式進(jìn)行存儲和讀取的，異常塊與正常塊往往具有不同的模式，因此可以通過減少異常重建的情況來提高檢測精度。此外，MemAE 方法只對最后一層輸出特征圖應(yīng)用記憶模塊，導(dǎo)致重建模糊的問題，本文方法通過skip connection 結(jié)構(gòu)減少了重建模糊的問題，進(jìn)而提升了檢測性能。DAAD+方法以統(tǒng)一尺寸的記憶塊存儲和讀取，但是異常的大小不一，難以選擇合適的記憶塊尺寸，若尺寸過大，則導(dǎo)致重建模糊問題，若尺寸過小，則異常塊很可能與正常塊有相同的模式，從而導(dǎo)致重建異常問題。而本文方法通過使用多個不同尺度的記憶塊來適應(yīng)不同大小的異常，既可以提升重建清晰度，又可以減少重建異常的情況，因此本文方法檢測精度高于DAAD+方法，在15 個類別中，該方法在11 個類別上取得了最高的AUC 分?jǐn)?shù)，并達(dá)到最高的平均AUC，比DAAD+方法高出1.5%。

3.4 消融實(shí)驗(yàn)

通過在MVTec AD 數(shù)據(jù)集上做消融實(shí)驗(yàn)，以驗(yàn)證所提出的各個模塊的有效性，實(shí)驗(yàn)結(jié)果如表2 所示。將GANomaly 生成網(wǎng)絡(luò)替換為UNet 并刪去skip connection 為基準(zhǔn)，對加入的模塊進(jìn)行消融實(shí)驗(yàn)。

表2 不同算法對模型性能的影響Table 2 The impact of different algorithms on model performance

分析表2 的數(shù)據(jù)可知，當(dāng)對基準(zhǔn)網(wǎng)絡(luò)直接加入skip connection 后，平均AUC 下降了4.3%，因?yàn)閟kip connection 雖然降低了網(wǎng)絡(luò)的重建難度，增強(qiáng)重建清晰度，但是對于異常圖像，異常特征從編碼器通過skip connection 直接傳入解碼器用于重建，導(dǎo)致輸入圖像與重建圖像差別不大，難以判斷是否存在異常，不同算法的檢測結(jié)果對比見圖4。如圖4（b）和圖4（c）所示，圖4（b）重建圖像較為模糊，圖4（c）重建圖像與原始圖像極為相似（其中，圖4（b）、圖4（c）和圖4（d）的第1 列表示預(yù)測圖像，第2 列表示預(yù)測圖像和其對應(yīng)的原始圖像的殘差圖像）。加入塊金字塔記憶模塊后，平均AUC 提升了14.2%，因?yàn)橥ㄟ^塊金字塔記憶模塊存儲正常圖像的典型特征，過濾掉大部分異常特征，避免了異常特征通過skip connection 傳入解碼器用于重建，減少了重建異常的情況，增大了異常圖像和其重建圖像之間的差異性，同時還提升了重建清晰度，如圖4（d）所示，重建圖像清晰，明顯改善重建異常問題。加入SSIM 損失，平均AUC 提升了0.4%，提升了網(wǎng)絡(luò)的重建效果。

圖4 不同算法的檢測結(jié)果對比Fig.4 Comparison of detection results of different algorithms

4 結(jié)束語

本文提出一種基于塊金字塔記憶模塊的無監(jiān)督異常檢測方法。針對重建模糊和重建異常等問題，設(shè)計塊金字塔記憶模塊用來存儲不同尺度的正樣本信息，在進(jìn)行異常檢測時，該模塊將異常特征轉(zhuǎn)化為正常特征，抑制異常重建，從而通過重建誤差來檢測異常，并且該模塊能夠聚合與輸入特征相似的不同尺度的記憶項(xiàng)，更好地抑制較大面積差的異常重建，提升細(xì)膩度重建能力。利用skip connection 結(jié)構(gòu)將經(jīng)過塊金字塔記憶模塊增強(qiáng)后的不同尺度的特征圖送入解碼器能夠提高重建的清晰度，并采用SSIM 損失函數(shù)增強(qiáng)原始圖像和重建圖像的差異性。在MVTec AD 數(shù)據(jù)集上的實(shí)驗(yàn)結(jié)果表明，本文方法相較于其他主流方法，在平均ROC 指標(biāo)上的得分有明顯提升，減少了異常重建情況和提升了重建清晰度。下一步將針對塊金字塔記憶模塊多個尺度的增強(qiáng)特征圖，探索更加有效的融合方式來提升模型的檢測精度，同時優(yōu)化模型結(jié)構(gòu)，加快推理速度。