基于知識蒸餾的心律失常分類模型

張 逸，周 莉，陳 杰

（1.中國科學院 微電子研究所，北京 100029；2.中國科學院大學，北京 100049）

隨著社會的不斷發(fā)展，人們的生活節(jié)奏越來越快，心血管疾病已成為導致人類死亡的主要疾病之一。心律失常是由心臟活動的起源或傳導障礙導致的異常心跳節(jié)律，是心血管疾病發(fā)病前的常見癥狀。如果能在心律失常發(fā)生的前期對其作出診斷，就能減少心血管疾病的發(fā)作，降低心血管疾病的死亡率。在臨床診斷中，醫(yī)生常用心電圖（Electrocardiogram，ECG）來檢測和診斷心律失常。ECG 是從體表記錄心臟每一心動周期所產(chǎn)生的電活動變化圖形的技術。ECG 心拍波形的變化可以作為判斷心律失常類型的有效依據(jù)。

隨著計算機技術和人工智能技術的不斷發(fā)展，ECG 自動分類技術發(fā)展迅速?？纱┐魇皆O備的廣泛普及也使得利用可穿戴設計對ECG 進行實時監(jiān)測成為現(xiàn)實。深度神經(jīng)網(wǎng)絡目前被應用在ECG 心律失常分類中，取得了較好的分類效果，同時，也有著計算量大、占用內(nèi)存多的缺點，難以部署在可穿戴設備等資源受限的硬件上。文中提出的基于知識蒸餾和注意力機制的輕量級神經(jīng)網(wǎng)絡模型，能在保持網(wǎng)絡的參數(shù)量和計算量較小的同時，具備良好的分類性能。

1 ECG心律失常分類

1.1 ECG心律失常分類模型

ECG 心律失常分類模型一般分為3 個步驟：ECG信號提取與預處理、ECG 信號特征提取、ECG 心律失常分類[1]。

圖1 ECG心律失常分類模型

在心律失常分類部分，已經(jīng)涌現(xiàn)出了許多ECG自動分類方法，主要分為基于傳統(tǒng)機器學習的方法和基于深度神經(jīng)網(wǎng)絡的方法。

1.2 基于傳統(tǒng)機器學習的方法

運用于ECG 心律失常分類的基于傳統(tǒng)機器學習的方法主要有線性分類器、支持向量機（Support Vector Machine，SVM）、K 近鄰算法等[2-6]?；趥鹘y(tǒng)機器學習的方法雖然能取得較好的分類效果，但仍具有一定的不足，這些方法需要人工進行特征提取，不能自動提取特征，分類效果依賴于模式空間的人為構(gòu)建，存在一定的局限性。

1.3 基于深度神經(jīng)網(wǎng)絡的方法

隨著人工智能技術的快速發(fā)展，人工神經(jīng)網(wǎng)絡中的深度神經(jīng)網(wǎng)絡近年來被廣泛應用在圖像處理和語音識別等多個研究領域中。卷積神經(jīng)網(wǎng)絡（Convolutional Neural Networks,CNN）目前被廣泛應用于心律失常分類中[7-8]。文獻[9]中提出了一種使用卷積神經(jīng)網(wǎng)絡對住院病人的ECG 信號進行分類的方法，取得了較好的分類效果。文獻[10]中提出使用二維卷積神經(jīng)網(wǎng)絡AlexNet 和VGG-16 對ECG 圖像進行分類。使用卷積神經(jīng)網(wǎng)絡進行ECG 分類不需要人工設計特征，將ECG 數(shù)據(jù)輸入后，CNN 中的卷積層可以自動提取ECG 數(shù)據(jù)的特征。深度卷積神經(jīng)網(wǎng)絡在網(wǎng)絡訓練過程中能充分挖掘ECG 特征間的關聯(lián)，具有良好的分類性能。

2 該文方法

該文構(gòu)建了輕量級二維卷積神經(jīng)網(wǎng)絡模型對ECG 心律失常進行分類。主要分為3 個階段，分別是ECG 輸入數(shù)據(jù)增強、基于注意力機制的輕量級神經(jīng)網(wǎng)絡模型設計、基于知識蒸餾的網(wǎng)絡優(yōu)化。

2.1 ECG輸入數(shù)據(jù)增強

該研究使用由美國麻省理工學院提供的研究心律失常的MIT-BIH 數(shù)據(jù)庫作為數(shù)據(jù)集。首先需要定義MIT-BIH 數(shù)據(jù)集每個ECG 心拍的范圍，并對其進行切片分成每個單獨的ECG心拍，將每個ECG心拍的范圍定義為以R 峰為基準點，取R 峰前100 個采樣點和R 峰后140 個采樣點作為一個心拍。經(jīng)過切片后，共生成約10 萬個包含正常類型和心律失常類型的ECG 心拍。按照美國醫(yī)療儀器促進協(xié)會(Association for the Advancement of Medical Instrumentation，AAMI)

對ECG 心拍分類的標準，這些心拍中約有90%的樣本為正常類別樣本，心律失常類別的樣本數(shù)量僅占10%。樣本數(shù)量的不平衡會導致訓練后的神經(jīng)網(wǎng)絡對多數(shù)類樣本過擬合，對少數(shù)類樣本的分類靈敏度較低。因此需要對ECG 輸入數(shù)據(jù)進行數(shù)據(jù)增強。該文使用合成少數(shù)類過采樣算法（Synthetic Minority Over-sampling Technique，SMOTE）[11]對一維ECG 信號進行數(shù)據(jù)增強。SMOTE 算法是基于隨機過采樣算法的一種改進方案，對數(shù)據(jù)中的少數(shù)類樣本進行分析并模擬出新樣本，之后將新樣本添加到數(shù)據(jù)集中。最后將每個一維ECG 心拍轉(zhuǎn)換為適合于輸入二維CNN 模型的二維圖像數(shù)據(jù)。

2.2 輕量級神經(jīng)網(wǎng)絡模型設計

為了降低網(wǎng)絡的計算量和復雜度，該文在滿足分類精度高的前提下，設計盡可能輕量化的二維CNN 模型。該模型由基于深度可分離卷積和注意力機制的基本模塊構(gòu)成。深度可分離卷積最初是在文獻[12]中提出。該卷積模式將標準卷積分成深度卷積（Depth-wise Convolution，DW）和1×1 的逐點卷積（Poin-wise Convolution，PW）兩步進行。在卷積核大小為3 時，保持相同的表達能力的前提下，深度可分離卷積的計算量僅為標準卷積的1/8 左右，可以大大減少卷積核參數(shù)數(shù)量和卷積計算的復雜度。注意力機制目前被廣泛應用在自然語言處理、圖像識別等深度神經(jīng)網(wǎng)絡模型中。在文獻[13]中將注意力機制用在計算機視覺領域。由于卷積運算通過將通道和空間信息融合在一起來提取特征，該文在每個基本模塊中的深度卷積之后加入了結(jié)合通道注意力和空間注意力的CBAM 模塊[14]，通過CBAM 模塊來強調(diào)通道和空間部分的有意義的特征。CBAM 模塊將注意力過程分成兩個獨立的部分，可以減少所需的參數(shù)量和計算量。CBAM 模塊示意圖如圖2 所示。

圖2 CBAM模塊示意圖

通道注意力模塊主要是探索不同通道間的feature map，示意圖如圖3 所示。卷積后的feature map 經(jīng)過MaxPool 池化和AvgPool 池化聚合特征圖的空間信息，生成最大池化特征和平均池化特征。這兩個特征經(jīng)過多層感知器MLP 后進行加和和sigmoid 激活后得到最終的通道注意力特征Mc。

圖3 通道注意力模塊示意圖

空間注意力模塊的示意圖如圖4 所示。通道注意力Mc和輸入的feature map 經(jīng)elementwise 相乘后生成空間注意力模塊的輸入特征圖。輸入特征圖沿通道作MaxPool 池化和AvgPool 池化后聚合feature map 中的通道信息，再利用一個卷積層對其進行卷積，經(jīng)sigmoid 激活后得到空間注意力特征Ms，最后將空間注意力特征Ms與原輸入的feature map 相乘得到新的feature map。

圖4 空間注意力模塊示意圖

網(wǎng)絡基礎模塊示意圖如圖5 所示。在每個模塊內(nèi)使用2 次1×1 的PW 卷積和1 次3×3 的DW 卷積，DW 卷積層后接CBAM 模塊。

圖5 網(wǎng)絡基礎模塊示意圖

該文提出的輕量級卷積神經(jīng)網(wǎng)絡結(jié)構(gòu)如表1 所示。輸入ECG 圖像后，接入普通卷積層，普通卷積層后接5 個網(wǎng)絡基礎模塊，基礎模塊后再加一層普通卷積。該網(wǎng)絡使用全局平均池化層（Global Average Pooling）加Softmax 層替換全連接層作為網(wǎng)絡的分類器，可以減少參數(shù)量和計算量。

表1 輕量級卷積神經(jīng)網(wǎng)絡結(jié)構(gòu)

2.3 基于知識蒸餾的網(wǎng)絡優(yōu)化

為了進一步提升網(wǎng)絡的分類精度，該文基于知識蒸餾[15]的思想，對網(wǎng)絡進行優(yōu)化。知識蒸餾是指用預先訓練好的一個模型復雜、層數(shù)較深的大型網(wǎng)絡作為教師網(wǎng)絡，將教師網(wǎng)絡學習到的信息和知識在輕量級的學生網(wǎng)絡訓練的過程中傳遞給它，從而能夠提高輕量級網(wǎng)絡的分類效果。所用的知識蒸餾方法是通過將教師網(wǎng)絡產(chǎn)生的分類概率分布作為軟知識來輔助硬知識訓練學生網(wǎng)絡。在學生網(wǎng)絡訓練時，損失函數(shù)設定為：

其中，Lhard為學生網(wǎng)絡的類別預測概率qi與真實標簽值ti之間的交叉熵損失。Lsoft為學生網(wǎng)絡的類別預測概率qi與教師網(wǎng)絡的輸出類別概率pi的交叉熵損失。

由于正確類別對應的真實標簽值ti為1，其余類別對應的為0，而教師網(wǎng)絡的輸出類別概率pi包含了各個類別的預測概率，這樣就導致Lhard中包含的信息量相比于Lsoft包含的信息量較少。使用知識蒸餾的方法可以將Lsoft中包含的信息輸入到學生網(wǎng)絡訓練的監(jiān)督信息中去，從而在不增加網(wǎng)絡的參數(shù)量和計算量的前提下提升網(wǎng)絡的分類精度。

3 實 驗

3.1 數(shù)據(jù)集生成

實驗使用由MIT-BIH 數(shù)據(jù)庫經(jīng)過數(shù)據(jù)增強后生成的ECG 圖片作為數(shù)據(jù)集。其中訓練集共79 006 張ECG 圖片，測試集共19 316 張ECG 圖片。

3.2 實驗結(jié)果

該文實驗使用基于MobileNetV2 設計的ECG 分類網(wǎng)絡作為知識蒸餾部分的教師網(wǎng)絡，學生網(wǎng)絡為設計的輕量級卷積神經(jīng)網(wǎng)絡。在ECG 圖像數(shù)據(jù)集上進行訓練和測試。基于AAMI 標準，該模型將ECG心拍類別分為4類，測試結(jié)果如表2和表3所示。

表2 ECG測試集上的混淆矩陣

表3 準確率、模型參數(shù)量和計算量比較

表2 給出了該模型在測試集上得到的混淆矩陣?；煜仃嚨牧写砹烁黝悇e的預測值，行代表了各類別的真實值，能夠表示各類別的真實值和預測值的分布情況。

表3 給出了該模型在測試集上的準確率、參數(shù)量和FLOPs，F(xiàn)LOPs 為浮點運算次數(shù)，可以反映模型的計算量大小。與其他方法相比，該文提出的模型參數(shù)量和計算量較小，分類準確率能夠達到98.3%。

4 結(jié)論

該文提出了一種基于知識蒸餾和注意力機制的心律失常分類模型，設計了輕量級二維卷積神經(jīng)網(wǎng)絡對ECG 心律失常進行分類。實驗表明，該模型網(wǎng)絡的參數(shù)量和計算量較小，在ECG 心律失常分類中達到了98.3%的分類準確率，取得了較好的分類效果。