基于NVAE和OB-Mix的小樣本數據增強方法

楊 瑋，鐘名鋒，，楊 根，侯至丞，王衛軍，袁 海

1.陜西科技大學 機電工程學院，西安 710021

2.廣州先進技術研究所 機器人與智能裝備中心，廣州 510000

3.中南大學 機電工程學院，長沙 410083

4.廣東技術師范大學 自動化學院，廣州 510000

近年來，計算機視覺的發展極其迅速，目標檢測作為計算機視覺的子課題也取得了突破性的進展[1-4]。然而，目標檢測的前沿性理論難以適用于工程實踐，一方面是因為工程實踐中對檢測實時性以及檢測準確性要求較高；另一方面，在工業領域中要獲取大量的標注數據需要投入較高成本，這成為了優秀的前沿理論應用于工程實踐中最大的障礙[5]。為解決模型對數據量的依賴問題，前人在數據生成及數據正則化策略方面做了大量的研究工作。

在數據生成方面，文獻[6]提出了一種人工缺陷仿真的方法來模擬手機屏幕四種經典的缺陷。對于線性缺陷，首先通過隨機游走算法生成缺陷骨架，再將骨架連接、擴張形成缺陷區域，最后通過高斯分布確定線性缺陷區域灰度值；而對于塊狀缺陷，先生成一條隨機的不規則閉合曲線，再采用不同的原理分別生成斑點區域和簇區域，最后根據高斯分布或者現有紋理來確定生成塊狀區域的灰度值。該文章提出的方法可以生成與真實缺陷高度相似且多樣性豐富的缺陷圖像，但是缺陷生成過程極其繁瑣復雜，還需要人為構造檢測目標的特征描述，可實現性較差。文獻[7]提出一種人工缺陷合成算法生成汽車內部零件缺陷圖像，實現了94%的檢測率。文獻[8]在金屬制品缺陷檢測系統中使用GAN來生成數據集，該研究表明生成數據量的增加可以提高缺陷識別準確率。文獻[9]在GAN的框架下提出一種具有編解碼結構的仿真網絡，該方法引入GAN的對抗訓練思想，僅需要少量的缺陷訓練數據即可生成特定形狀和類型的模擬缺陷。文獻[10]提出了一種基于深度學習的豆類缺陷檢測方案，并結合了一種GAN結構的自動標記數據增強方法，通過該方案訓練得到的檢測模型準確率達到80%。然而，基于GAN數據生成的方法生成器和判別器在對抗訓練時容易失衡[11-12]，訓練過程不穩定。并且此類方法均是使用整圖進行對抗訓練生成數據，引入大量不必要信息，使得網絡更加難以收斂。

在mix數據增強方面，為了進一步提高卷積神經網絡（CNN）在各種機器視覺任務上的表現，研究人員提出了許多網絡訓練的正則化策略[13-15]。特別地，為了防止CNN過度關注檢測目標對象局部的明顯特征，文獻[16]提出了CutMix數據增強策略對訓練圖像中檢測對象的局部進行剪切，再用另一幅圖像的補丁替換移除區域，標簽也按組合圖像的像素數比例混合。從而使網絡關注整個對象區域，提高模型的泛化能力和定位能力。文獻[17]提出mixup方法從訓練數據中隨機抽取兩個示例進行線性融合，該方法實際是通過結合特征向量的線性插值會導致相關目標的線性插值的先驗知識來擴展訓練分布，實驗表明mixup方法有效提高模型的魯棒性。文獻[18]采樣兩個未標記的示例A和B，將兩者分別輸入分割網絡，得到結果SA和SB，再對A和B、SA和SB進行相應的像素級融合，生成新的帶像素標簽的虛擬數據。該方法在Cityscapes和Pascal VOC 2012上進行性能評估，結果表明該方法有效改進現有技術。以上mix數據增強的工作是利用兩個訓練實例進行不同形式的mix策略融合生成新的樣本，事實證明mix數據增強的正則化策略對于模型訓練是有良性增益的。

因此，本文提出了一種先生成檢測目標對象圖像，然后將目標對象圖像與正常樣本背景圖像融合的數據擴充方法（如圖1所示，pre_pro、enc_tower、dec_tower、post_pro為NVAE的網絡模塊（在1.2.1小節中詳細說明）；?和θ為網絡的權重參數；μ和σ2為多元高斯混合模型的均值和方差；ε為變分自編碼器的隨機性參數；z為自編碼器的潛變量），以滿足在工業領域深度學習對大量數據的需求。主要創新點為利用數據分布擬合的方式挖掘圖像特征信息，同時提出生成目標與背景隨機數量、位置融合的網絡訓練正則化策略，并且提出了基于上述兩點的數據增強方法。本文的具體工作有：首先，在少量已標注樣本中，根據標簽信息分離目標對象和背景，引入新派變分自編碼器（nouveau variational auto-encoder，NVAE）[19]構建檢測目標的數據分布模型，再通過采樣潛變量的方式生成與真實目標屬于同一分布的全新目標圖像；然后，在得到生成目標圖像后，本文提出了OB-Mix（object-background mixture）均值移植的圖像融合方法，以實現將生成目標圖像與無缺陷背景圖像進行隨機位置融合，從而構建出新的合成圖像數據；最后利用合成的圖像數據進行目標檢測模型YOLOv5[20]的訓練，本方法在僅使用474張標注圖像以及400張無檢測目標的背景圖像的情況下，使YOLOv5的檢測精確率達到95.86%。

1 方法

1.1 研究對象

本文采用某電子紙膜片的麻點缺陷圖像作為研究對象。如圖2所示，麻點缺陷是由許多細小的黑點組成，分布在電子紙的任意位置，僅當黑點達到一定數量影響人的視覺感受時才認為是缺陷。如果采用傳統圖像處理方法檢測，很難設計檢測特征及檢測標準。深度學習具有較強的語義信息提取能力，能夠模擬人的感官進行缺陷的識別與檢測。但深度學習依托于大量樣本的訓練，在工業領域中，很難滿足海量缺陷數據的需求。

圖2 電子紙缺陷Fig.2 Electronic paper defects

針對難以獲取大量缺陷產品數據集的問題，本文提出了一種基于NVAE數據生成和OB-Mix圖像融合的數據擴充方法，通過本方法在小樣本的情況下盡可能地提高深度學習模型的檢測準確率。小樣本訓練模型的關鍵在于充分利用已有樣本的信息，盡量減少冗余信息的干擾，所以本方法提出將目標對象和背景進行分離，如圖3所示。在檢測模型訓練中，模型只需要識別目標對象的特征，并不需要過于關注背景[21]。分離目標和背景，一方面可以更有針對性地挖掘目標對象的特征信息，也有利于NVAE擬合分布訓練時模型的收斂，另一方面可以引入mix正則化策略，將生成的不同缺陷目標對象與不同的無缺陷的背景圖像進行OB-Mix隨機位置融合，以此增加數據集的背景豐富性以及目標對象位置的任意性，提高模型的定位能力和泛化能力。

圖3 分離缺陷和背景Fig.3 Separating defects and background

1.2 基于NVAE的目標圖像生成

NVAE是一種改進的變分自編碼器（variational autoencoder，VAE）[22]網絡。NVAE通過編碼將一個高維變量（缺陷圖像）映射到低維的特征變量（潛變量），再通過解碼重構另一個全新的高維變量。網絡的訓練目標就是通過輸入變量與輸出變量不斷編碼映射和解碼重構來學習到缺陷圖像分布與潛變量分布之間的分布轉換關系，從而可以采樣潛變量的分布來生成缺陷數據。NVAE的核心思想是圖像數據的本質是分布模型，該模型控制著數據特征的生成，而目標檢測網絡是通過特征提取骨干提取出圖像數據的特征，然后利用該特征去檢測目標對象。通過擬合檢測目標數據的分布模型，即可以生成在此分布下具有相同特征的目標對象圖像，從而實現了樣本數據集的擴充。

1.2.1 NVAE理論解析

對于給定高維變量x,x是由真實的條件分布控制生成的，對應存在一個隱藏的連續隨機低維變量z,x表征屬性由z的各維控制，z是由先驗概率分布生成。生成數據的目標是求出，根據式（1），采樣來生成x數據。

式（4）中，L(?,?*;x)為變分下界（evidence lower bound，ELBO）。E表示期望。從自編碼器的角度來看，等式右邊第一項為重建損失（reconstruction loss），第二項為正則化損失（regularization loss）。對這兩項進行梯度估計進而優化模型，即可實現最小化真實分布和擬合分布的KL散度的目標。

文獻[22]提出SGVB（generic stochastic gradient variation Bayes）估計器進行梯度估計，對于重建損失項的優化目標是使得給定z重建出的x與真實x的概率盡量大，可以通過MSE的方式實現，如式（5）：

其中，x為由z重構的生成圖像，c為正則化系數，取值為x的維數。而對于正則化損失項，可以寫成式（6）：

其中，μ、σ2為多元高斯混合模型的均值與方差。

同時使用了重參數化方法（reparameterization trick）將數據的隨機性歸結到參數ε中，如式（7）：

NVAE為提高擬合的z先驗概率分布q?(z)表達能力以及q?(z|x)的構建能力，將潛變量分組為z={z1,z2,…,zN}，令：

其中，z

NVAE在具體實施如圖4所示，使用兩個神經網絡分別作為擬合的分布模型，通過網絡訓練的方式求參數?和θ的最優解。編碼器（encoder）網絡輸出的是特征屬性分布的均值μ與方差的對數lbσ2（假設近似的后驗分布q?(z|x)為具有對角協方差結構的多元高斯混合模型，如式（10）），均值與方差用于計算后驗分布q?(z|x)與假設的潛變量分布q(z)的KL散度以及采樣生成潛變量z，同時z會加入隨機項ε，如式（11）：

圖4 NVAE模型結構Fig.4 Model structure of NVAE

其中，N(μ,σΙ)表示具有對角協方差結構的多元高斯混合模型，⊙表示一種運算方式，在NVAE網絡中通過Concat操作實現。

得到z之后將其輸入解碼器（decoder）重建出模擬的缺陷圖像x，再計算出真實缺陷圖像x與重構的均方誤差MSE，將MSE以及上述的負KL散度作為網絡模型的損失，進行反向傳播優化。NVAE網絡的具體實現網絡結構如圖5和表1。圖5中，x為真實圖像，z為低維特征，ε為變分自編碼器的隨機性參數，為生成圖像。

表1 NVAE網絡結構模塊說明Table 1 Modules description of NVAE structure

圖5 NVAE網絡結構模塊Fig.5 Structure modules of NVAE network

1.2.2 NVAE檢測目標圖像生成

在數據生成時使用474張帶缺陷及相應標簽的高分辨率圖像，從中分離521張缺陷圖像；400張無缺陷圖像。用521張缺陷圖像訓練NVAE網絡，設置400個迭代輪次（epoch），批數量（batch size）為16，一個迭代輪次分為24個迭代步數（step）。網絡擬合10組潛變量z的分布，通過考慮在給定x時，每組潛變量z的擬合分布與假設服從的分布間的KL散度來評價其擬合效果，每組潛變量KL散度的計算方式如式（12），每一組潛變量的KL散度結果及總的KL散度結果如圖6。

圖6 潛變量的分布擬合結果Fig.6 Distribution fitting results of latent variables

其中，q?(z(i,l)|x)為已知缺陷數據x時網絡構建z服從的分布，q(z(i,l))為假設服從的分布，和Δσ2均由卷積網絡直接產生。

在圖6中，分別顯示了每組潛變量z與所有z總的KL散度的優化過程。其中每一組z擬合的分布與假設的分布之間的KL散度都已優化趨于穩定（z9還在收斂）。而所有潛變量z綜合的KL散度也趨于最小值24.71，并且達到穩定，這說明訓練的潛變量z的分布已經與假設的分布已到達最小距離。需要說明的是，本研究中兩種分布模型最小距離的優化并沒有一個通用的衡量標準，本實驗的訓練優化結果僅表明在本研究條件下的可行解空間中找到相對最優的一個結果。

對于網絡構建q?(z|x)、pθ(x|z)以近似真實分布的綜合效果，采用ELBO來評價，其計算公式如式（4），訓練結果如圖7。ELBO是由重建損失和先驗正則項兩項組成，優化重建損失目的是給定z后重建出的與真實x相似的概率盡可能大；而優化先驗正則項是為了讓擬合后驗分布逼近潛變量的假設分布，同時讓后驗分布具有一定的隨機性；因此ELBO是一種“Bias-Variance”平衡。

圖7 NVAE的ELBO優化結果Fig.7 ELBO of NVAE

由圖7可以看到，NVAE網絡的優化在第5 000 step之后已基本收斂穩定，ELBO在值為2 764.42達到平衡，這也表明NVAE網絡在擬合缺陷數據分布時其缺陷圖像的重構能力以及重構圖像的隨機性在該規則下最優。

NVAE在完成400迭代輪次的訓練之后，采樣潛變量z進行圖像生成，生成效果如圖8。

1.3 背景自適應OB-Mix數據增強

本文基于mix正則化策略對網絡訓練的增益作用，提出一種全新的OB-Mix正則化策略，OB-Mix是一種將NVAE生成的檢測對象與背景進行融合的策略。在目標檢測任務中，網絡應該關注檢測目標對象的特征共性，背景只作為與目標的對比。為此，本文提出將生成目標對象在背景圖像上隨機位置、數量融合的方式。以屬于同一分布下的缺陷圖像與不同背景圖像融合得到的數據，可以體現目標對象的特征共性以及背景的隨機性和豐富性，使網絡學習到更多目標對象的特征信息。這種增加數據集的背景豐富性以及目標對象位置的任意性的方式，是利于提高模型對目標對象的識別和定位能力的。

真實數據中缺陷的尺寸通常是不同的，為了提高網絡模型對不同缺陷尺寸的魯棒性，在進行OB-Mix前先對NVAE生成的缺陷圖像進行尺度變換操作，圖9是真實數據的缺陷尺寸統計數據，由圖可以知道缺陷的寬是在20～60 pix之間，而高是在20～70 pix之間，寬高均主要集中在30～50 pix之間。圖中顏色深度表示缺陷數量，顏色越深數量越大。故在融合時會將缺陷圖像寬高在20～60 pix區間內進行隨機變換。

圖9 缺陷尺寸統計結果Fig.9 Statistic of defect size

根據經驗，通常每張電子紙膜片上出現1～2個缺陷的情況較多，出現數量更多的缺陷概率較小。故本文引入Beta分布(α=2,β=5)擬合目標對象在每一張電子紙圖像中的數量分布情況，分布模型如圖10，缺陷數量K∈(0,4]?Z。同時在圖像融合時使用了灰度值均值移植的方法，即將生成的缺陷圖像要替換區域的灰度值均值移植到缺陷圖像上，使其能更好地與背景融合。融合后的效果如圖11。

圖10 服從Beta分布的缺陷數量概率Fig.10 Probability of number of defects obeying Beta distribution

圖11 真實數據與合成數據對比Fig.11 Comparison of real data and synthetic data

2 實驗與結果分析

為驗證本文提出的小樣本數據增強方法對檢測網絡的訓練效果，采用YOLOv5作為驗證模型。訓練數據包括帶標簽的真實數據474張圖像，利用474張真實圖像以及400張無缺陷的背景圖像生成的合成數據400張。測試數據為264張真實圖像。設置兩組對照實驗數據，一組是474張真實圖像加400張無缺陷背景圖像，另一組是474張真實圖像加400張生成圖像。使用兩組數據訓練網絡進行對比，訓練epoch為300代，每組數據分別進行5次訓練實驗，取5次結果平均值。網絡訓練的檢測精確率及召回率對比結果如圖12。

圖12 不同訓練數據的模型精確率和召回率Fig.12 Precision and recall of different training data

從圖12中可以看到，在訓練中加入了400張生成圖像數據后，網絡的檢測精確率及召回率都比僅有真實數據時高。從訓練的結果數據來看，使用真實數據加生成數據訓練的精確率在92%左右浮動，最高達到95.52%；而只用真實數據訓練的精確率在85%左右浮動，最高為85.33%。對于召回率，在使用了生成數據時其值在89%左右浮動，最高達到90.02%；僅使用真實數據時其值在81%左右浮動，最高達到82.98%。另一個重要的點是使用了生成數據訓練網絡使其收斂速度更快，且收斂過程也更加穩定。以上結果說明本文提出的基于生成數據的小樣本訓練方法是有效的，對于本實驗條件下的YOLOv5檢測精確率提高11.94個百分點，檢測召回率提高8.48個百分點，并且收斂速度提高，收斂過程更穩定。

為進一步研究本文提出方法的效果，對YOLOv5網絡訓練過程的損失收斂情況進行分析，邊界框定位損失和置信度損失訓練結果如圖13。

圖13 不同訓練數據的邊界框損失和置信度損失Fig.13 Bounding box loss and confidence loss of different data

對于圖（a）的定位損失，加入了生成數據訓練后，可以更快地收斂到一個更低的值。而對于圖（b）置信度損失，僅使用真實數據訓練的收斂情況在訓練前期要優于加入生成數據的。在50 epoch之后兩者的收斂情況幾乎一致，但是使用了生成數據的損失在隨著epoch的增加越來越貼近僅使用生成數據的損失（局部放大圖）。

在本文提出的方法中，除了使用數據分布模型擬合以生成數據來解決小樣本問題的思想之外，還有另一個重要的機制，即mix機制，從圖13中（a）的結果可以提出猜想，OB-Mix對網絡的目標定位能力有提升作用。為了探究這一猜想，進行了另一組實驗：對同一張背景圖像，分別設置不同的缺陷圖像OB-Mix融合數量，觀察網絡邊界框定位損失的收斂情況以及mAP_0.5：0.95的指標情況。實驗時依然假設每張電子紙圖像中缺陷數量服從Beta分布(α=2,β=5)，在此分布基礎上加上一個基數以實現不同缺陷數量等級，分別設置基數為0，2，4，實驗結果如圖14。

圖14 不同mix缺陷數量對網絡定位性能的影響Fig.14 Effect of different mixture defect numbers on localization performance of YOLOv5

如圖14（a）顯示，缺陷圖像mix基數在0～4的范圍內，隨著mix數量的增多，網絡的邊界框損失可以收斂到更低的值，且收斂速度更快。而對于圖14（b），當缺陷的mix數量增多時，可以大幅地提高mAP_0.5：0.95的值，并且在mix基數為4時，訓練時網絡的定位性能的提高過程也更穩定。以上結果表明，在0～4的mix數量范圍內，使用OB-Mix方法將更多數量的缺陷圖像融合到背景圖像上，可以有效提高YOLOv5的定位性能。

如圖15為不同mix基數下訓練的YOLOv5網絡檢測精確率。圖15表明mix缺陷基數從0增加到4時，網絡的檢測精確率也在不斷提升，并且相較于mix基數為0時，mix基數為2和4訓練過程中，網絡的精確率提升也更加穩定，但是mix基數從2增加到4時，精確率增幅較為微小。

圖15 不同mix缺陷數量對網絡精確率的影響Fig.15 Effect of different mixture defect numbers on precision of YOLOv5

最后，使用測試集對不同訓練集訓練出的模型進行測試，測試結果如表2所示。使用了生成數據（mix基數為0）訓練的模型精確率提高了11.73個百分點，召回率提高了4.21個百分點。當生成圖像的mix基數從0增加到4時，精確率和召回率都有相應的提升。在mix基數為4時，精確率提高了17.60個百分點，召回率提高了18.28個百分點。

表2 不同訓練集訓練的模型精確率與召回率Table 2 Precision and recall under different training data 單位：%

不同數據訓練出來的模型對電子紙表面缺陷進行檢測的檢測結果如圖16。圖中紫色帶標簽方框為模型檢測結果，標簽顯示的是預測的置信度閾值；黃色方框為漏檢缺陷；紅色方框為誤檢缺陷。從圖中可以看到，mix缺陷基數從0增加到4時，漏檢缺陷數量在減少，網絡預測的置信度閾值在提高，這是由于網絡定位性能的提高進而使得IoU（intersection over union）值提高。而誤檢則是因為該區域出現與麻點缺陷相似的斷斷續續劃痕缺陷，導致模型誤判。

圖16 不同訓練集訓練的模型檢測結果Fig.16 Detection results trained on different datasets

使用不同的目標檢測模型對提出的數據增強方法進行驗證，同樣采用上述的對比實驗數據組（生成數據mix基數為4），分別訓練各個模型，統計5次訓練結果數據的平均值如表3，其中AP表示average precision。DLP（decrease in localization performance）表示IoU閾值從0.5增加到0.75時，AP下降的值，DLP可以反映模型的定位精度變化情況。由表3可以看出，使用本數據增強方法的模型AP值均有不同程度的提升，可以證明本方法對目標檢測模型的性能提升具有普適性。在IoU閾值為0.5情況下，使用本方法擴充400張圖像數據可以對Faster-RCNN、Retinanet、SSD和YOLOv5的AP值分別提升了25.73%、22.63%、22.2%、42.68%。對于四個目標檢測模型，使用生成數據訓練的模型DLP值均比僅使用真實數據的低，說明本方法對四個模型的定位性能均有提升，其中對于YOLOv5模型的提升幅度最大。

表3 不同訓練數據下目標檢測模型AP值Table 3 AP of object detection models under different data 單位：%

為進一步研究本方法對網絡性能的提升效果，將本方法與其他常用的數據增強方法進行對比實驗，實驗結果如圖17所示。

圖17 常用數據增強方法訓練的網絡精確率Fig.17 Precision of network trained by different data augmentation methods

由圖17，none為不使用任何數據增強方法的結果。在常用的數據增強方法中，僅有隨機裁剪（crop）方法起到增益作用，精確率提高了3.52個百分點。旋轉（rotate）方法僅僅使得網絡收斂更快，精確率上沒有任何提升，其他兩種數據增強方法（flip和crop&flip&rotate）反而使精確率有所下降。相比于傳統常用的方法，本文提出的方法在收斂速度和精確率上都有大幅度的優勢。

以上所有實驗結果表明，在數據樣本不足情況下，本文提出的方法可以有效提高目標檢測網絡的檢測性能，其中YOLOv5尤為明顯，并且提升效果大幅優于目前常用的數據增強方法。在僅使用474張缺陷圖像的情況下，利用本方法可以將YOLOv5的檢測精確率提高到95.86%（mix基數-4），相比于不使用本方法的78.26%，提高了17.60個百分點。本研究的另一個貢獻是提出OB-Mix正則化策略并對其進行研究，該策略對目標檢測網絡的定位性能有增益作用，當采用本方法用小樣本進行檢測網絡訓練時，可以在一定范圍內進一步增加缺陷的mix數量來提高網絡的性能。

3 結束語

本文提出了一種基于NVAE數據生成和OB-Mix圖像融合的小樣本數據增強方法，用以解決在有限數據情況下訓練的目標檢測網絡準確率低的問題。首先從帶標簽的樣本數據中分離出缺陷目標圖像，再用NVAE去擬合缺陷圖像的分布模型，NVAE擬合是通過神經網絡構造一個缺陷圖像和其潛變量的分布模型轉換關系，然后假設潛變量服從一種簡單的分布，通過變分理論間接求出缺陷圖像的近似分布，這樣即可采樣生成全新的缺陷圖像。將生成的缺陷圖像通過OB-Mix方法融合到背景圖像中得到合成數據。本文將合成的數據結合真實數據訓練YOLOv5網絡，使得網絡在僅有474張缺陷圖像以及400背景圖像的情況下，可以達到95.86%的檢測精確率，相比于不使用本文方法的結果提高了17.60個百分點。進一步實驗證明，本方法對其他的目標檢測模型也同樣適用。

本方法的核心是在有限的數據情況下，充分挖掘檢測對象的本質信息，對于其他領域也同樣具有很大的應用潛力。未來研究工作可將本方法應用于不同領域圖像的檢測或分割等任務。另外，可以進一步探究mix正則化策略對網絡其他性能的影響。