一行命令跑評測,港中文MMLab開源自監督表征學習代碼庫OpenSelfSup
自監督表征學習發展迅速,但也存在諸多問題。近日,香港中文大學多媒體實驗室(MMLab)和南洋理工大學的研究者開源了一套統一的自監督學習代碼庫 OpenSelfSup。
前言
近幾個月來自監督表征學習領域獲得了顯著突破,特別是隨著 Rotation Prediction、DeepCluster、MoCo、SimCLR 等簡單有效的方法的誕生,自監督表征學習大有超越有監督表征學習的趨勢。
然而,做這個領域的研究者都深有感觸:1)自監督任務復雜而多樣,不同方法各有各的專用訓練代碼,難以結合、復用和改進;2)評價方案不統一,不同的方法難以在公平環境下對比;3)動輒百萬千萬的訓練數據量,訓練效率是個大問題。
針對這些問題,香港中文大學多媒體實驗室(MMLab)和南洋理工大學的研究者最近開源了一套統一的自監督學習代碼庫:OpenSelfSup。
開源地址:https://github.com/open-mmlab/OpenSelfSup
OpenSelfSup
統一的代碼框架和模塊化設計
OpenSelfSup 使用 PyTorch 實現,支持基于分類、重建、聚類、memory bank、contrastive learning 的多種自監督學習框架,目前收錄了 Relative Location、Rotation Prediction、DeepCluster、OnlineDeepCluster、NPID、MoCo、SimCLR 等一系列表現較好的自監督表征學習方法,后續還將陸續跟進學術界最新算法。
在這個框架下,每個算法被拆解為 backbone、neck、head、memory_bank (optional)、hook (optional) 等多個可獨立設計的模塊,每個模塊均提供多個可選方案,開發者也可以自行設計各個模塊。
SimCLR 中的 backbone、neck 和 head 設計
標準化的評測方案
OpenSelfSup 目前支持 ImageNet/Place205 Linear Classification、ImageNet Semi-Supervised Classification、PASCAL VOC07 Linear SVM、PASCAL VOC / COCO Object Detection 等多個基準的評測方案,可以做到一行命令跑評測,非常方便。
高效率的分布式訓練
OpenSelfSup 收錄的算法全部實現了多機多卡的分布式訓練、提特征和測試。
上手容易
該代碼庫的環境配置、數據配置均有 from scratch 的腳本或詳細指導,簡單易行。訓練和測試現有算法,只需要一行命令即可搞定。
高度靈活性和可擴展性
1. 改進現有算法。OpenSelfSup 用 config 文件來定義各種參數和模塊,方便修改。config 中還支持一些較復雜的調整,比如 data augmentation 的組合、learning rate schedule、獨立調整某些網絡參數的優化參數等。例如,你希望單獨調整 head 中 fully-connected layer 的 momentum 和 learning rate,或者 backbone 中某幾層的 weight decay 等等,那么你可以在 config 中 optimizer: paramwise_option 下用正則表達式篩選出對應網絡參數,然后指定這些值,而不需要改動代碼。
DeepCluster中指定head(fc layer) 的momentum為0
2. 使用自己的訓練數據。OpenSelfSup 將數據處理和數據源解耦。使用自己的訓練數據,只需要新增一個自己的 data source 文件,并在 config 中指定即可。
3. 設計自己的算法。高度模塊化的設計使得自行設計算法變得非常簡單。你只需要設計自己的模塊或者復用現有模塊,然后新建一個模型文件將這些模塊組織起來即可。如果你的算法包含間隔 iteration、間隔 epoch 的操作(例如 DeepCluster 需要在每個 epoch 結束后對訓練集做聚類),也只需要新建一個 hook 來定義這些操作,并在 config 中設置這個 hook 即可。hook 的調用是自動進行的。
DeepCluster 初始聚類和每 n 個 epoch 后做聚類
結語
OpenSelfSup 是一個長期維護的開源項目,旨在方便學術界和工業屆的研究者推動自監督表征學習領域繼續前進。也希望有興趣的研究者、開發者幫助繼續完善 OpenSelfSup,為這個領域貢獻自己的力量。