InsightFace: 用OneFlow輕鬆實現超大規模人臉識別模型

人臉識別技術可以準確識別出圖像中的人臉和身份，具有豐富的應用場景，譬如金融場景下的刷臉支付、安防場景下的罪犯識別和醫學場景下的新冠流行病學調查等等。人臉識別的算法演變經歷了以 PCA 為代表的早期階段，再到以「人工特徵+分類器」為主的統計學習方法階段，近幾年，隨著大數據及 GPU 算力的爆發，人臉識別進入到深度學習算法為絕對主角的階段。

人臉識別的規模也從最初基於幾百張圖片、識別數十張人臉，到現在動輒識別上百萬甚至億級別的人臉，這樣超大規模的人臉識別需求，帶來了諸多要解決的技術難題，特別是基於分布式機器學習系統訓練大規模人臉識別模型的挑戰。採用混合併行的方式解決超大規模人臉識別問題，已經成為工業界的最佳實踐，但主流深度學習框架（如TensorFlow, PyTorch, MXNet等）僅支持了更容易實現的數據並行，而如果要實現混合併行，往往需要基於深度學習框架進行二次開發且編程複雜度高。OneFlow 作為主打分布式易用和性能的框架，有著去中心化的Actor機制及SBP的抽象，天然對分布式有著良好支持，訓練速度快、顯存佔用低，同樣的顯卡能支持更大的batch size，從單機多卡拓展到多機分布式訓練並不需要複雜的設置。本文基於OneFlow實現了大規模人臉識別的解決方案，該方案可以幫助用戶輕鬆使用數據並行 + 模型並行的混合訓練，同時還支持了Partial FC採樣技術，理論上支持過億(ID數)級別的人臉識別。InsightFace [1]是基於 MXNet 框架實現的業界主流人臉識別解決方案。相較MXNet的實現方案，基於OneFlow的實現方案在性能方面更是十分優秀，OneFlow在數據並行時速度是其2.82倍；模型並行時速度是其2.45倍；混合併行+Partial fc時速度是其1.38倍。基於OneFlow實現的代碼已合併至 insightface的官方倉庫[2]，其中包含了數據集製作教程、訓練和驗證腳本、預訓練模型[3]以及和MXNet模型的轉換工具。在 The 1:1 verification accuracy on InsightFace Recognition Test (IFRT) 驗證集上，Oneflow及MXNet訓練模型的精度對比如下：FrameworkAfricanCaucasianIndianAsianAll
https://github.com/deepinsight/insightface/tree/master/recognition/oneflow_face[4]下面，將以大規模人臉識別任務為例，介紹基於OneFlow和MXNet框架是如何現業界流行的大規模人臉識別技術，以及數據並行、模型並行以及Partial FC採樣的實現細節，主要內容包含：完全理解本文的技術內容需要對 OneFlow 的並行觀和原理有所了解，OneFlow團隊曾寫過幾篇文章詳細介紹了OneFlow的並行觀、SBP屬性等概念，有需要的讀者可參考：1.大規模人臉識別背景介紹前面的概述簡單介紹了人臉識別的一些應用場景以及超大規模人臉識別所帶來的一些挑戰，下面將更具體一點，來分析下大規模人臉識別方案中涉及到的關鍵技術以及具體難在何處。1.1 面臨的問題簡單的網絡基於深度學習網絡的大規模人臉識別方法一般都是基於常規的 CNN 網絡（如resnet50、resnet100等）來提取輸入圖片中的人臉特徵，然後人臉特徵會輸入到全連接層，最後基於全連接的輸出計算 Loss，單純從網絡結構來看，大規模人臉識別相關的網絡通常如下圖所示：

複雜的MarginLoss大規模人臉識別面臨的第一個問題是，當使用常規的CNN+FC的網絡，後接softmax交叉熵損失函數來對人臉進行分類時，往往得到的結果並不能令人滿意，同一個人在不同角度，不同光線等情況下，往往會得到不一樣的分類結果。所以在最後計算交叉熵損失之前，需要通過一些方法對FC層的輸出進行處理，譬如加入特徵間距離的度量，用於判斷類別內部（同一個人的不同臉部）和外部（不同人臉間）之間的差異大小。簡單來說，一個魯棒的人臉分類模型應該可以使得映射後的特徵具有較小的類內和較大的類間距離。通過這種基於類間、類內距離度量的方法衍生出了一系列距離度量損失函數，如 center loss、triplet loss、sphereface loss、cosface loss、arcface loss 等等。其中arcface loss源於論文《ArcFace: Additive Angular Margin Loss for Deep Face Recognition》[8]，由於應用了超球面度量、餘弦距離等創新方法，其對於類間和類內距離度量更為精準，在各種人臉識別任務和數據集上表現優異。arcface.png在 insightface 的源碼中，對應 arcface loss 和 cosface loss 的實現是 MarginLoss[9] 類：

class MarginLoss(object):
    """ Default is Arcface loss
    """
    def __init__(self, margins=(1.0, 0.5, 0.0), loss_s=64, embedding_size=512):
        """
        """
        # margins
        self.loss_m1 = margins[0]
        self.loss_m2 = margins[1]
        self.loss_m3 = margins[2]
        self.loss_s = loss_s
        self.embedding_size = embedding_size

    def forward(self, data, weight, mapping_label, depth):
        """
        """
        with autograd.record():
            norm_data = nd.L2Normalization(data)
            norm_weight = nd.L2Normalization(weight)
            #
            fc7 = nd.dot(norm_data, norm_weight, transpose_b=True)
            #
            mapping_label_onehot = mx.nd.one_hot(indices=mapping_label,
                                                 depth=depth,
                                                 on_value=1.0,
                                                 off_value=0.0)
            # cosface
            if self.loss_m1 == 1.0 and self.loss_m2 == 0.0:
                _one_hot = mapping_label_onehot * self.loss_m3
                fc7 = fc7 - _one_hot
            else:
                fc7_onehot = fc7 * mapping_label_onehot
                cos_t = fc7_onehot
                t = nd.arccos(cos_t)
                if self.loss_m1 != 1.0:
                    t = t * self.loss_m1
                if self.loss_m2 != 0.0:
                    t = t + self.loss_m2
                margin_cos = nd.cos(t)
                if self.loss_m3 != 0.0:
                    margin_cos = margin_cos - self.loss_m3
                margin_fc7 = margin_cos
                margin_fc7_onehot = margin_fc7 * mapping_label_onehot
                diff = margin_fc7_onehot - fc7_onehot
                fc7 = fc7 + diff

            fc7 = fc7 * self.loss_s
            return fc7, mapping_label_onehot
MarginLoss包含m1~m3這3個參數，通過這3個參數的組合來實現cosface loss和arcface loss。而在oneflow中的實現中則更為簡便，直接調用flow.combined_margin_loss[10]即可實現MarginLoss系列的功能。巨大的人臉ID數導致顯存爆炸對於工業界的人臉識別業務，人臉 ID 數通常會超過百萬級，甚至可以達到千萬級至億級別，在這種情況下全連接層的參數矩陣通常會超出單個 GPU 設備的顯存上限，所以，僅僅靠普通的數據並行也無法完成訓練。而這就是大規模人臉識別方案的核心挑戰所在。1.2 解決方案數據並行 or 模型並行為了處理上面提到的問題，工業界對於超大規模的人臉識別任務，往往採用數據並行 + 模型並行的混合併行方式。即在網絡前面的CNN部分，採用數據並行進行人臉特徵提取，而最後的全連接層則採用模型並行，將參數矩陣切分到多個 GPU 上。2.OneFlow 如何實現大規模人臉識別基於OneFlow實現的大規模人臉識別方案對齊了 insightface官方的partail_fc[11] 的實現（基於MNXet），支持數據並行、數據|模型混合併行和Partial FC採樣技術，在loss方面支持設置了 m1,m2和m3超參以定義 softmax loss、arcface loss、cosface loss 以及其他組合形式的 combined loss。代碼已合併至insightface官方倉庫—oneflow_face[12]。下面將通過整體結構和技術細節實現這兩個層面來介紹基於OneFlow的大規模人臉識別方案。2.1 整體結構首先是 採用數據並行的CNN 特徵提取部分，CNN提取的特徵(Features)作為後面的全連接(FC)層的輸入，全連接層採用模型並行。全連接層fc1經過Margin loss layer(fc7)處理後的輸出，同label一起計算softmax交叉熵損失，得到最終的loss。
圖中展示了每個GPU設備上具體的計算流程，在GPU上方有全連接層的權重矩陣Weight（圖中的matmul節點），黃色長方體表示的Features經CNN提後取的人臉特徵。對於batch_size大小的批量圖片輸入，Feature的形狀為 (batch_size, emb_size)，emb_size根據網絡不同通常為128或512。圖中的權重(Weight)的大小與人臉類別 ID 數有關，在大規模人臉識別的工業實踐中，類別 ID 數通常為百萬到億級別，假設類別 ID 數為1千萬，則模型大小為(emb_size, 10000000)。經過全連接層的特徵矩陣和權重矩陣相乘後((batch_size, emb_size) × (emb_size, 10000000))的輸出特徵形狀為 (batch_size, 10000000)。在OneFlow的實現方案中全連接層採用模型並行，即對權重矩陣做切分而使用全量的特徵數據，因此輸入特徵的 SBP 屬性為 Broadcast，即表示每個GPU設備上都會拷貝一份特徵數據；而權值的 SBP 屬性為 Split(1)，即參數矩陣在維度1被切分到各個 GPU 設備上，假設有 P 個 GPU，則每個 GPU 上有 (emb_size, 10000000/P) 大小的權值 。全連接層的輸出形狀取決於輸入features的形狀以及類別ID數，故每個設備上的輸出形狀為 (batch_size, 10000000/P)，且 SBP 屬性也為 Split(1)。由於本方案對權值做了切割(Split(1))，故通常來說，需要對全連接層的輸出做合併，並轉為按 Split(0) 切分的數據並行，但是由於全連接層的輸出數據塊較大，如果直接由 Split(1) 轉為 常規的Split(0)，會引入大量(可以避免的)通信。因此，在 OneFlow 算子實現的內部，並不先進行 Split 的轉化，而是將全連接層的輸出直接作為 softmax 的輸入進行計算，因為 softmax 的運算特性，可以使得輸出的 SBP 屬性依然是 Split(1)。類似的，softmax 的輸出（按 Split(1) 切分）繼續作為 sparse_cross_entropy 的輸入進行計算，由於算子本身的特性和 OneFlow 的機制，sparse_cross_entropy 的輸出的 SBP 屬性依然可以保持 Split(1)。經過 sparse_cross_entropy 處理後的輸出，邏輯上獲得最終的 loss 結果。此時，數據塊的形狀為 (batch_size, 1)，已經很小，這時候再將模型並行的Split(1）模式轉為按 Split(0) 切分的數據並行。以上即是 OneFlow 實現的全連接層的內部工作流程，對於普通算法開發者來說，了解以上內容即掌握了OneFlow大規模人臉方案中全連接層處理的核心流程。對於框架開發者和實現細節感興趣的朋友，請看下面的小節—2.全連接層的技術細節。這一小結，將會用較大篇幅展開softmax和sparse_cross_entropy實現細節相關的內容以及Split(1）切分是如果做到數學上的等價。2.2 Oneflow實現代碼解析下面，我們講解一下在Oneflow是如何通過簡單的幾行代碼來實現大規模人臉識別方案。首先，backbone部分的網絡是類似的由CNN+FC全連接層構成，我們重點看一下FC之後的Marginloss層及相關處理。主要代碼如下：
elif config.loss_name == "margin_softmax":
            if args.model_parallel:
                print("Training is using model parallelism now.")
                labels = labels.with_distribute(flow.distribute.broadcast())
                fc1_distribute = flow.distribute.broadcast()
                fc7_data_distribute = flow.distribute.split(1)
                fc7_model_distribute = flow.distribute.split(0)
            else:
                fc1_distribute = flow.distribute.split(0)
                fc7_data_distribute = flow.distribute.split(0)
                fc7_model_distribute = flow.distribute.broadcast()
            fc7_weight = flow.get_variable(
                name="fc7-weight",
                shape=(config.num_classes, embedding.shape[1]),
                dtype=embedding.dtype,
                initializer=_get_initializer(),
                regularizer=None,
                trainable=trainable,
                model_name="weight",
                distribute=fc7_model_distribute,
            )
            if args.partial_fc and args.model_parallel:
                print(
                    "Training is using model parallelism and optimized by partial_fc now."
                )
                (
                    mapped_label,
                    sampled_label,
                    sampled_weight,
                ) = flow.distributed_partial_fc_sample(
                    weight=fc7_weight, label=labels, num_sample=args.total_num_sample,
                )
                labels = mapped_label
                fc7_weight = sampled_weight
            fc7_weight = flow.math.l2_normalize(
                input=fc7_weight, axis=1, epsilon=1e-10)
            fc1 = flow.math.l2_normalize(
                input=embedding, axis=1, epsilon=1e-10)
            fc7 = flow.matmul(
                a=fc1.with_distribute(fc1_distribute), b=fc7_weight, transpose_b=True
            )
            fc7 = fc7.with_distribute(fc7_data_distribute)
            fc7 = (
                flow.combined_margin_loss(
                    fc7, labels, m1=config.loss_m1, m2=config.loss_m2, m3=config.loss_m3
                )
                * config.loss_s
            )
            fc7 = fc7.with_distribute(fc7_data_distribute)
        else:
            raise NotImplementedError

        loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
            labels, fc7, name="softmax_loss"
        )

        lr_scheduler = flow.optimizer.PiecewiseScalingScheduler(
            base_lr=args.lr,
            boundaries=args.lr_steps,
            scale=args.scales,
            warmup=None
        )
        flow.optimizer.SGDW(lr_scheduler,
                            momentum=args.momentum if args.momentum > 0 else None,
                            weight_decay=args.weight_decay
                            ).minimize(loss)

        return loss
和MXNet的實現類似，fc1表示backbone網絡最後的全連接層輸出；fc7表示後面的Marginloss層。在模型並行的情況下，需要分別設置數據、fc1、fc7層的SBP屬性：
fc1_distribute = flow.distribute.broadcast()     #  表示SBP為B，廣播，用於全量數據的同步

fc7_model_distribute = flow.distribute.split(0)  #  表示SBP為S0，表示模型並行且按第0維切割
# （FC層為dense層，在框架實現上會對dense層模型做轉置，model_distribute設Split(0)代表模型在第1維分割，SBP屬性相當於Split(1)

fc7_data_distribute = flow.distribute.split(1)   #  表示SBP為S1，表示數據並行且按第1維切割（和模型並行保持一致）
設置SBP屬性後，Oneflow框架會根據其SBP屬性及內在的Boxing機制，在後續前向反向的過程自動完成模型切分、數據的同步、以及內在的調用集合通信源語來完成broadcast、allreduce相關操作。接下來，獲取fc7層的權重矩陣fc7_weight，其SBP形狀為fc7_model_distribute設定的模型切分。接著會判斷是否使用Partial fc採樣，如果使用，則通過flow.distributed_partial_fc_sample()獲取採樣後的權重及label。1.首先，通過flow.matmul()進行矩陣乘法(fc1層->fc7層的前向傳播)，並設置其結果的SBP屬性為fc7_data_distribute。2.然後，通過flow.combined_margin_loss()完成fc7層的前向傳播，由於其SBP屬性並沒有改變，故繼續設置其SBP屬性為fc7_data_distribute。3.最後，直接調用flow.nn.sparse_softmax_cross_entropy_with_logits()即可完成求softmax交叉墒損失，至此完成了整個過程前向+loss計算，之後的反向過程、梯度更新、學習率更新等將由後面的flow.optimizer自動完成。2.3 模型並行的技術細節上文已經介紹了OneFlow 大規模人臉方案的整體實現，實現中全連接層採用模型並行，即全部的數據與部分的模型進行 matmul 運算，其原理示意可參考 Consistent 與 Mirrored[13] 一文中的相關部分，之後對matmul的結果進行 softmax+求交叉熵損失函數，即調用flow.nn.sparse_softmax_cross_entropy_with_logits()完成整個計算過程。可能，有的讀者會好奇，模型並行下的softmax是如何通過一個sparse_softmax_cross_entropy_with_logits計算得到的？模型並行下的計算和正常情況下的計算數學上等價嗎？這一小節，將詳細介紹在FC層模型並行下的softmax交叉熵計算過程，以及其數學上為何等價。問題產生的原因首先，正常|數據並行情況下的softmax過程要求對全局數據進行softmax，即對label和完整的特徵(模型)進行數學計算，但在模型並行的實現方案中，由於全連接層為模型並行，故softmax的計算方式和正常方案有所不同。在OneFlow的模型並行方案中，label與部分的模型進行 matmul 運算，然後直接在各卡本地計算 softmax，最終通過 sparse cross entropy 算子計算得到 SBP 屬性為 Split(1) 的 loss，最終將 loss 的 SBP 屬性轉為 Split(0)， 即轉為和普通數據並行方案等價的loss。因此，看起來和普通的softmax方案會存在差異，那麼這就帶來了這兩種方案下的softmax計算在數學上是否等價的問題！下文介紹此問題產生的原因以及 OneFlow 的 flow.nn.parse_softmax_cross_entropy_with_logits 算子，是如何在模型並行的情況下做到數學上softmax計算等價的。模型並行下的softmax在邏輯上，全連接層後，應該進行 softmax 運算，其公式如下：在實際的深度學習框架實現中，為避免數值溢出，往往會調整以上公式，讓指數部分減去數據中的 max 值。也就是說，實際採用的公式為：此時，由於本方案採用模型並行，每個設備上只有部分的模型，如果在每張卡本地直接套用以上 softmax 公式，就會出現問題：公式中涉及到的求 max 和求 sum 的運算均需要全局的信息如果在求 max 和求 sum時忽略掉全局信息，只對每張卡獨立地、使用本地數據進行 softmax 計算，得到的結果與數學上的要求不一致。但是，如果將全局信息，也就是全連接層的輸出廣播到各個 GPU 設備上，假定類別 ID 數為10000000，則數據塊的大小為 (batch_size, 10000000)，通信代價較高，因此如何高效地分布式完成 softmax 計算也常是其它框架在實現模型並行時面臨的難題之一。下面，將介紹 OneFlow 的算子 flow.nn.parse_softmax_cross_entropy_with_logits 的內部實現原理和細節，可以看到它利用 OneFlow 框架 SBP 機制，是如何巧妙地解決以上問題。sparse_softmax_cross_entropy_with_logits 的實現細節為了解決以上產生的分布式計算softmax的問題，本方案將邏輯上的全局 softmax，在網絡內部拆分成多個Operator(op)進行計算/操作，在經過reduce max、sub、exp...等一系列op計算後，最終得到全局 softmax 的等價計算結果，以下是具體實現的示意圖：reduce max的作用是按列取出最大，以形狀為 (3, 3) 的數據塊為例，其效果為：
[5, 4, 2]                 [5]
[4, 3, 3]  -reduce max->  [4]
[1, 9, 8]                 [9]
這樣，在每個卡上經過獨立的 reduce max 運算後，得到 Partial max 的結果，即每個 GPU 設備上是部分結果。此時，數據的塊的形狀已經變小為 (bathc_size, 1)，為了將多個 GPU 上的 Partial max 結果求 max 得到全局的max結果，需要將每張卡上的部分max結果廣播到其他卡，這樣每張卡上都能獲取到全局max結果，在這一步驟中雖然仍然不可避免地需要 AllReduce 通信，但是因為 reduce max 後的數據塊已經大大減小，因此降低了通信成本。接著，可以進行邏輯上的減法操作，即圖中的 Sub。Sub 的輸入有兩個：來自全連接層的輸出（SBP 屬性為 Split(1)），以及上一步計算得到的全局 max 的結果（SBP 屬性為 Broadcast），OneFlow 可以自動推導出，Sub 後的輸出結果的 SBP 屬性為 Split(1)。然後，減法的輸出，經過圖中的 Exp 運算，完成了原數學公式中的分子的求解。需要多說明一句的是，Sub 及 Exp 的過程中，都保持了 Split(1) 的性質，因此數據塊小，效率高。再之後，需要繼續求解原數學公式中的分母，此時需要使用 reduce sum 運算，同上文介紹的 reduce max 類似，reduce sum 計算取得的結果其實只是單獨卡上的局部結果，即 Partial sum,後續需要一次 Broadcast 通信並計算，才能得到全局 sum 的結果。得到數學公式中的分子分母后，使用 Div 算子得到他們相除的結果，Div 的輸入為之前步驟計算得到的數學公式中的分子與分母，其中分子的 SBP 屬性為 Split(1)，分母的 SBP 屬性為 Broadcast。OneFlow 可以自動推導出，相除的結果的 SBP 屬性為 Split(1)。可以看到以上過程，利用多個算子在 OneFlow 框架下的運算，既做到了與數學邏輯上的要求一致，又很大程度地降低了多個設備之間的通信量，因為 Partial max 及 Partial sum 過程中，通信的數據塊大小均為 (batch_size, 1)。接著求最終loss，從上文的介紹可知，softmax 的輸出的 SBP 屬性為 Split(1)，需要繼續經過 sparse cross entropy 計算得到 loss。sparse cross entropy 的輸入有兩個：前一層 softmax 的輸出，SBP 屬性為 Split(1)OneFlow 會根據 sparse cross entropy 輸入的 SBP 屬性，自動推導出其輸出的 SBP 屬性為 Partial sum。sparse_cross_entropy 的輸出，經過 Partial sum 後，再廣播到各個卡，同 softmax 的情況類似，因為數據塊已經經過減小，因此大大減少了通信量。OneFlow 中的代碼實現上文詳細討論了人臉識別網絡的全連接層，及其後續網絡在 OneFlow 中如何用一系列算子(op)進行實現。從理解原理的角度出發，以上討論的篇幅較長，但對於普通的分布式訓練用戶而言，使用 OneFlow 實現大規模人臉方案的代碼卻極為簡潔：
labels = flow.parallel_cast(labels, distribute =  flow.distribute.broadcast())
embedding = flow.parallel_cast(embedding, distribute = flow.distribute.broadcast())

fc7 = flow.layers.dense(
            inputs=embedding,
            units=args.class_num,
            model_distribute=flow.distribute.split(0),
) #dense中模型會做轉置，model_distribute設Split(0)代表模型在第1維分割，SBP屬性為Split(1)

loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
        labels, fc7.with_distribute(flow.distribute.split(1)), name="softmax_loss"
)
在以上代碼中，先使用 flow.parallel_cast 方法將 labels 和 embedding 的 SBP 屬性設置為 Broadcast。然後，在flow.layers.dense內通過設置其 model_distribute 參數為 flow.distribute.split(0) 將模型的 SBP 屬性設置為 Split(1)，從而創建了模型並行方式的全連接層。最後，通過調用 flow.nn.sparse_softmax_cross_entropy_with_logits 獲取 loss。上文討論原理中所涉及的 SBP 類型推導、SBP 屬性轉換等工作，都由 OneFlow 框架自行完成。3.MXNet的大規模人臉識別方案基於MXNet實現的大規模人臉方案面臨的主要問題主要有：混合併行模式下常規的算子如gloss.SoftmaxCrossEntropyLoss 無法正常工作的問題。MXNet的解決方案為採用horovod、類似numpy的mxnet.ndarray對矩陣進行手動切分計算和手寫softmax交叉熵相關的代碼，整體實現相對複雜。下面我們將對MXNet實現方案的實現流程做簡單講解，包含前向、反向和Marginloss相關部分。3.1 MXNet實現代碼解析模型訓練入口是 train_memory.py 中的 train_module.fit()[14]：
train_module.fit(train_data_iter,
                     optimizer_params=backbone_kwargs,
                     initializer=mx.init.Normal(0.1),
                     batch_end_callback=call_back_fn)
通過調用 SampleDistributeModule[15] 類中定義的 fit 方法來開啟整個訓練過程。在 fit 方法中除了會對模型的參數做初始化外，還主要包括以下內容：通過 forward_backward（）[16]來進行前向、反向的過程，並計算出待更新的梯度；通過 update()[17] 調用 optimizer 完成對模型、學習率等參數的更新。下面將重點講解forward_backward（）[18]中前向、反向這兩個部分的具體流程。
def forward_backward(self, data_batch):
        """A convenient function that calls both ``forward`` and ``backward``.
        """
        total_feature, total_label = self.forward(data_batch, is_train=True)
        self.backward_all(total_feature, total_label)
前向
def forward(self, data_batch, is_train=None):
        self.backbone_module.forward(data_batch, is_train=is_train)
        if is_train:
            self.num_update += 1
            fc1 = self.backbone_module.get_outputs()[0]
            label = data_batch.label[0]

            total_features = self.allgather(tensor=fc1,
                                            name='total_feature',
                                            shape=(self.batch_size * self.size,
                                                   self.embedding_size),
                                            dtype='float32',
                                            context=self.gpu)
            total_labels = self.allgather(tensor=label,
                                          name='total_label',
                                          shape=(self.batch_size *
                                                 self.size, ),
                                          dtype='int32',
                                          context=self.cpu)
            return total_features, total_labels
        else:
            return None
其中 backbone_module 主要是提取圖片特徵，如果不添加Marginloss層，則人臉識別網絡的架構是由一個CNN接FC層組成，然後通過softmax交叉熵計算損失函數。如果添加了Marginloss層則還需要對Marginloss做相關的處理，在MXNet的實現中將Marginloss層的前向和反向過程一起放在了下面的backward_all()函數中。反向由前向部分的描述可知，反向部分包括兩個主要部分：1.Marginloss層的處理2.普通CNN網絡的反向下面，我們看一下backward_all()方法:
 def backward_all(
        self,
        total_feature,
        total_label,
    ):
        # get memory bank learning rate
        self.memory_lr = self.memory_optimizer.lr_scheduler(self.num_update)

        self.grad_cache = self.get_ndarray(self.gpu, 'grad_cache',
                                           total_feature.shape)
        self.loss_cache = self.get_ndarray(self.gpu, 'loss_cache', [1])

        self.grad_cache[:] = 0
        self.loss_cache[:] = 0

        if not bool(config.sample_ratio - 1):
            grad, loss = self.backward(total_feature, total_label)
        else:
            grad, loss = self.backward_sample(total_feature, total_label)

        self.loss_cache[0] = loss

        total_feature_grad = grad
        total_feature_grad = hvd.allreduce(total_feature_grad, average=False)

        fc1_grad = total_feature_grad[self.batch_size *
                                      self.rank:self.batch_size * self.rank +
                                      self.batch_size]
        self.backbone_module.backward(out_grads=[fc1_grad / self.size])
第1部分，對Marginloss層的處理集中放在了backward()和backward_sample()部分，二者區別在於backward()對應的是sample_ratio=1.0時的反向（不使用Partial fc進行採樣）；backward_sample()對應使用Partial fc進行採樣時的反向過程。第2部分，實現CNN網絡的反向，即：self.backbone_module.backward通過這兩個部分，完成了整個反向過程的梯度計算。下面重點看一下第1部分對Marginloss層的處理。Marginloss
def backward(self, total_feature, label):
        memory_bank = self.memory_bank
        assert memory_bank.num_local == memory_bank.num_sample, "pass"

        _data = self.get_ndarray2(self.gpu, "data_%d" % self.rank,
                                  total_feature)
        # Attach grad
        _data.attach_grad()
        memory_bank.weight.attach_grad()

        # Convert label
        _label = self.get_ndarray2(self.gpu, 'label_%d' % self.rank, label)
        _label = _label - int(self.rank * memory_bank.num_local)
        _fc7, _one_hot = self.fc7_model.forward(_data,
                                                memory_bank.weight,
                                                mapping_label=_label,
                                                depth=memory_bank.num_local)

        # Sync max
        max_fc7 = nd.max(_fc7, axis=1, keepdims=True)
        max_fc7 = nd.reshape(max_fc7, -1)

        total_max_fc7 = self.get_ndarray(context=self.gpu,
                                         name='total_max_fc7',
                                         shape=(max_fc7.shape[0], self.size),
                                         dtype='float32')
        total_max_fc7[:] = 0
        total_max_fc7[:, self.rank] = max_fc7
        hvd.allreduce_(total_max_fc7, average=False)

        global_max_fc7 = self.get_ndarray(context=self.gpu,
                                          name='global_max_fc7',
                                          shape=(max_fc7.shape[0], 1),
                                          dtype='float32')
        nd.max(total_max_fc7, axis=1, keepdims=True, out=global_max_fc7)

        # Calculate exp(logits)
        _fc7_grad = nd.broadcast_sub(_fc7, global_max_fc7)
        _fc7_grad = nd.exp(_fc7_grad)

        # Calculate sum
        sum_fc7 = nd.sum(_fc7_grad, axis=1, keepdims=True)
        global_sum_fc7 = hvd.allreduce(sum_fc7, average=False)

        # Calculate prob
        _fc7_grad = nd.broadcast_div(_fc7_grad, global_sum_fc7)

        # Calculate loss
        tmp = _fc7_grad * _one_hot
        tmp = nd.sum(tmp, axis=1, keepdims=True)
        tmp = self.get_ndarray2(self.gpu, 'ctx_loss', tmp)
        tmp = hvd.allreduce(tmp, average=False)
        global_loss = -nd.mean(nd.log(tmp + 1e-30))

        # Calculate fc7 grad
        _fc7_grad = _fc7_grad - _one_hot

        # Backward
        _fc7.backward(out_grad=_fc7_grad)

        # Update center
        _weight_grad = memory_bank.weight.grad
        self.memory_optimizer.update(weight=memory_bank.weight,
                                     grad=_weight_grad,
                                     state=memory_bank.weight_mom,
                                     learning_rate=self.memory_lr)

        return _data.grad, global_loss
其中fc7_model即MarginLoss[19]所在的層，細節如論文所示，這裡就不展開了。backwark()方法裡首先通過fc7_model.forward()完成將開始Marginloss層的前向過程並得到以及one_hot的label:
_fc7, _one_hot = self.fc7_model.forward(_data,
                                        memory_bank.weight,
                                        mapping_label=_label,
                                        depth=memory_bank.num_local)
然後計算softmax交叉墒損失，再通過反向產生梯度。通過調用 gloss.SoftmaxCrossEntropyLoss 即可完成softmax交叉墒損失的計算，不過對於數據+模型的混合併行的情況，MXnet 現有的api無法支持，需要手動實現整個softmax交叉墒的計算。具體來說就是通過下面一系列max、sum、div的計算再加上 allreduce 和 broadcast 集合通信操作共同完成。反向傳播完成後，再通過self.memory_optimizer.update完成optimizer裡模型權重的更新、學習率的更新。通過以上代碼分析，可以看出基於 MXNet實現的人臉識別方案還是比較複雜的，除了要求算法開發者對整個人臉識別流程、模型數據的切分和softmax交叉熵數學計算比較熟悉，還需要對分布式集合通信原理、horovod的使用較為熟練。整體來看，要求還是比較高的。4.數據並行、模型並行解決方案的通信量對比這一小結將介紹傳統的數據並行方案、以及本方案(數據+模型並行)中涉及到的通信量做大致分析及對比。4.1 數據並行通信量首先，若採用純數據並行，假定人臉類別 ID 數為10000000，涉及的通信量為大模型的反向梯度的 Allreduce，數據塊大小為 (emb_size, 10000000)， 若採用 RingAllReduce，總傳輸量為：2 * emb_size * 10000000 * (P - 1)4.2 混合併行通信量讓我們對純數據並行與本文的混合併行兩種解決方案的通信量進行比較，採用本文的混合併行，涉及的通信量由以下幾部分組成：1.CNN 網絡得到的人臉特徵由 Split(0)->Broadcast 的傳輸，數據塊大小為 (batch_size, emb_size)2.label 由 Split(0)->Broadcast 的傳輸，數據塊大小為 (batch_size, 1)3.Softmax 計算過程中，每卡 Split(1) 切分數據的 max 值由 Partial max->Broadcast 的傳輸，數據塊大小為(batch_size, 1)；每卡 Split(1) 切分數據的 sum 值由 Partial sum->Broadcast 的傳輸，數據塊大小為(batch_size, 1)4.最終算出 loss 後的通信，數據塊大小為 (batch_size, 1)(batch_size * emb_size + 6 * batch_size) * (P - 1)4.3 對比總結不難算出，當batch_size大小及顯卡數P一定、且人臉類別 ID 數巨大時(譬如1000萬)，純數據並行和基於本方案的混合併行，二者通信量存在數量級的差異：2 * emb_size * 10000000 * (P - 1)VSbatch_size * （emb_size + 6) * (P - 1)採用數據並行時，巨大的人臉類別 ID 數導致的顯存佔用可能撐爆顯存使得無法訓練，即使可以訓練，巨大的通信量也會使得吞吐率降低，訓練速度變慢。而採用OneFlow混合併行方案時，不僅可以大大降低顯存佔用，提升訓練速度，此外假定 GPU 數目為 P，則在連接層每個 GPU 上的模型是總模型大小的1 / P，可以通過擴展GPU設備數，支持超大規模的類別 ID 數。5.Partial FC採樣技術模型並行，可以解決巨大的人臉分類ID帶來的權重存儲和通信問題，因為無論類別多大，總可以通過擴展GPU設備數來解決。然而，在矩陣乘法的實現中，除了模型權重矩陣之外，全連接層fc1輸出的logits矩陣也同樣需要存放在GPU顯存中：
fc7 = flow.matmul(
   a=fc1.with_distribute(fc1_distribute), b=fc7_weight, transpose_b=True
)

在模型並行的方案中，當模型並行到P臺設備上時，每卡的logits尺寸為：
(batch_size_per_device * P, num_classes/P)
當num_classes逐漸增加時，需要增加設備數量P以支持訓練，但同時logits矩陣的第一維也會逐漸增加導致更高的顯存佔用，從而實際支持的最大人臉ID類別數也有上限，即不能通過無限拓展設備數P來支持更大規模的ID類別數。為了解決此問題，格靈深瞳在論文《Partial FC: Training 10 Million Identities on a Single Machine》[20]提出了Partial FC的採樣技術，簡單來說，Partial fc即對權重和logits矩陣的一種採樣，通過設置相應的採樣率(sample ratial)來達到節省內存，能支持更大類別數的效果。
根據論文中的描述，softmax函數中的負類在人臉表徵學習中的重要性並沒有那麼高，即無需用每個輸出特徵和完整的模型權重矩陣相乘（全採樣），可以通過sample ratial設置採樣率，譬如sample ratial=0.1則表示只採樣10%的權重矩陣，而經過採樣後不損失精度，通過此方式可以輕鬆支持更大類別數。insightface相關的代碼在官方倉庫：insightface/partial_fc[21]。
OneFlow中也支持了Partial FC採樣的功能，只需調用flow.distributed_partial_fc_sample算子即可對權重矩陣、標籤label進行採樣。使用Partial FC採樣的主要代碼如下：
(mapped_label, sampled_label, sampled_weight) = flow.distributed_partial_fc_sample(
        weight=fc7_weight, label=labels, num_sample= num_sample)
6.OneFlow和MXNet實現的性能對比我們在相同硬體環境下，測試了基於MXNet和OneFlow框架實現的大規模人臉識別方案，從吞吐率(速度)、支持的最大batch size、支持的最大人臉ID數規模等指標對兩個方案進行了對比。總體來說，基於OneFlow實現的大規模人臉方案，在單機、多機情況下的表現均大幅由於MXNet實現，具體表現在：1.在相同batch size下，吞吐率更高，訓練速度更快且多卡時，性能損失較小，更接近線性加速比。其中：數據並行時，1836.8 vs 650.8(samples/s)，速度是MXNet的2.82倍模型並行時，1854.15 vs 756.96(samples/s)，速度是MXNet的2.45倍混合併行+Partial fc時(4機32GPU)，6931.6 vs 5008.7(samples/s)，速度是MXNet的1.38倍，且加速比28.1 vs 22.82.對GPU顯存的管理水平和利用率更高，在相同硬體配置下（gpu顯存固定時）支持跑更大的batch size模型並行時，115 vs 96，較MXNet提升20%混合併行+Partial fc時，115 vs 96，較基於MXNet實現的Partial fc提升20%混合併行+Partial fc時，單機支持的最大人臉ID數(num classes)：1350萬 vs 1200萬，較基於MXNet的Partial fc提升12.5%6.1 數據並行6.2 模型並行6.3 混合併行 + Partial FC(sample ratio = 0.1)emore數據集glint360k數據集6.4 Max batch size per device以下數據為單機8卡、fp32精度模式下測得，旨在比較相同硬體環境下，不同框架對GPU顯存的管理利用能力，以訓練時所能支持的最大batch size大為優。
OneFlowMXNetmax batch size per devicemax batch size per device6.5 Max num classes以下數據為單機、fp16混合精度模式下測得，旨在比較相同硬體環境下，不同框架所能支持的最大人臉ID類別數(num classes)，以比較框架的性能邊界。modenode_numgpu_num_per_nodemax num classes(OneFlow)max num classes(MXNet)所有數據、測試報告及代碼見DLPerf倉庫：https://github.com/Oneflow-Inc/DLPerf#insightface[22]7.總結隨著深度學習及GPU算力的爆發，以深度學習算法為基礎的人臉識別方案已經得到廣泛應用，人臉識別方案相關的深度學習網絡通常以CNN為主，其網絡結構通常較為簡單，但難點在於超大規模人臉模型的訓練上。通常由於類別數非常大，導致模型大小通常超出單張顯卡的顯存上限，所以不得不使用數據並行+模型並行的方式來進行訓練，而各大框架對此類需求支持的並不是很好，工業/企業落地時往往需要需要算法工程師、框架開發工程師對各大深度框架進行二次開發或深度定製，實現起來往往較為複雜且訓練效率得不到保障。通過上文與MXNet實現的性能對比可以看出，使用OneFlow 實現的（數據+模型混合併行）超大規模人臉識別方案，其優勢有：
OneFlow復現、調試Insightface的過程中，需要特別感謝Insightface項目的發起人過佳以及Partial fc的作者格靈深瞳的安翔。首先感謝兩位提供了如此高效的大規模人臉識別方案，其次，在OneFlow方案實現過程中，他們給出了耐心細緻的指導，在數據集、測試方面也給予了大力支持，衷心感謝！
[1]
InsightFace : https://github.com/deepinsight/insightface
[2]
insightface的官方倉庫: https://github.com/deepinsight/insightface/tree/master/recognition/oneflow_face
[3]
預訓練模型: https://github.com/Oneflow-Inc/oneflow_face#Pretrained-model
[4]
https://github.com/deepinsight/insightface/tree/master/recognition/oneflow_face: https://github.com/deepinsight/insightface/tree/master/recognition/oneflow_face
[5]
《OneFlow 的並行特色》: https://docs.oneflow.org/extended_topics/model_mixed_parallel.html
[6]
《如何評價 7 月 31 日一流科技開源的深度學習框架 OneFlow？》: https://www.zhihu.com/question/409036335/answer/1373468192
[7]
《僅此一文讓您掌握OneFlow框架的系統設計（上篇）》: https://zhuanlan.zhihu.com/p/337851255
[8]
《ArcFace: Additive Angular Margin Loss for Deep Face Recognition》: https://arxiv.org/pdf/1801.07698.pdf
[9]
MarginLoss: https://github.com/deepinsight/insightface/blob/79aacd2bb3323fa50a125b828bb1656166604487/recognition/partial_fc/mxnet/memory_softmax.py
[10]
flow.combined_margin_loss: https://github.com/Oneflow-Inc/oneflow_face/blob/master/insightface_train.py#L358
[11]
insightface官方的partail_fc: https://github.com/Oneflow-Inc/DLPerf/tree/master/MxNet/InsightFace/PartailFC
[12]
oneflow_face: https://github.com/deepinsight/insightface/tree/master/recognition/oneflow_face
[13]
Consistent 與 Mirrored: https://docs.oneflow.org/extended_topics/consistent_mirrored.html#_3
[14]
train_module.fit(): https://github.com/deepinsight/insightface/blob/863a7ea9ea0c0355d63c17e3c24e1373ed6bec55/recognition/partial_fc/mxnet/train_memory.py#L158
[15]
SampleDistributeModule: https://github.com/deepinsight/insightface/blob/863a7ea9ea0c0355d63c17e3c24e1373ed6bec55/recognition/partial_fc/mxnet/memory_module.py#L14
[16]
forward_backward（）: https://github.com/deepinsight/insightface/blob/863a7ea9ea0c0355d63c17e3c24e1373ed6bec55/recognition/partial_fc/mxnet/memory_module.py#L118
[17]
update(): https://github.com/deepinsight/insightface/blob/863a7ea9ea0c0355d63c17e3c24e1373ed6bec55/recognition/partial_fc/mxnet/memory_module.py#L119
[18]
forward_backward（）: https://github.com/deepinsight/insightface/blob/863a7ea9ea0c0355d63c17e3c24e1373ed6bec55/recognition/partial_fc/mxnet/memory_module.py#L118
[19]
MarginLoss: https://github.com/deepinsight/insightface/blob/863a7ea9ea0c0355d63c17e3c24e1373ed6bec55/recognition/partial_fc/mxnet/memory_softmax.py#L6
[20]
《Partial FC: Training 10 Million Identities on a Single Machine》: https://arxiv.org/abs/2010.05222
[21]
官方倉庫：insightface/partial_fc: https://github.com/deepinsight/insightface/tree/master/recognition/partial_fc
[22]
https://github.com/Oneflow-Inc/DLPerf#insightface: https://github.com/Oneflow-Inc/DLPerf#insightface