摘要
我們提出了 SegGPT,一個在上下文中執行“分割一切”的通用模型。我們將各種分割任務統一為一個通用的上下文學習框架,通過將不同類型的分割數據轉換為相同的圖像格式來兼容這些差異。SegGPT 的訓練被建模為一個帶有隨機顏色映射的上下文著色問題。訓練目標是根據上下文完成多樣的任務,而非依賴特定的顏色。訓練完成后,SegGPT 可通過上下文推理(in-context inference)執行圖像或視頻中的任意分割任務,如實例、區域、部件、輪廓和文本等。SegGPT 在多種任務上進行了評估,包括小樣本語義分割、視頻目標分割、語義分割和全景分割。實驗結果顯示其在域內和跨域目標分割中都展現出強大的能力,無論是定性還是定量評估。
1. 引言
分割是計算機視覺中最基本的問題之一,其目標是在像素級上定位并重組有意義的概念,如前景、類別、目標實例等。近年來,針對各類分割任務(如前景分割 [41]、交互式分割 [51, 34]、語義分割 [32, 28, 54, 39]、實例分割 [18, 11, 2, 48] 和全景分割 [23, 5, 8])的算法取得了顯著進展,表現越來越精確且快速。
然而,這些專用的分割模型通常局限于特定的任務、類別、粒度和數據類型。當面臨不同場景(如分割新概念,或從圖像遷移到視頻)時,往往需要重新訓練一個新模型。這不僅代價高昂,還難以適應大量不同分割任務的需求,缺乏可持續性。
在本研究中,我們的目標是訓練一個單一模型,能夠處理多樣且無限的分割任務。這一目標面臨兩個主要挑戰:(1) 如何在訓練中融合截然不同的數據類型,如部件、語義、實例、全景、人物、醫學圖像、航空圖像等;(2) 如何設計一種可泛化的訓練機制,區別于傳統的多任務學習方式,具備任務定義的靈活性,并能處理跨領域任務(out-of-domain tasks)。
為了解決這些挑戰,我們提出了 SegGPT——一個在上下文中執行“分割一切”的通用模型。我們將分割視為一種通用的視覺感知格式,并將不同的分割任務統一到一個通用的上下文學習框架中 [46]。該框架通過將各種分割數據轉換為相同的圖像格式,使模型可以處理不同類型的數據。SegGPT 的訓練被建模為一個上下文著色問題,對每個樣本使用隨機顏色映射。模型的目標是根據上下文對類別、實例、部件等區域進行著色。通過使用隨機顏色,模型被迫依賴上下文信息來完成任務,而非依賴預定義顏色,這使訓練方式更具靈活性與泛化能力。其余訓練流程與 [46] 保持一致,使用基礎 ViT [42] 和簡單的 smooth-?1 損失函數 [17]。
訓練完成后,SegGPT 能夠在圖像或視頻中通過少量示例進行上下文推理,完成多種分割任務,如目標實例、區域、部件、輪廓、文本等。為了更有效地整合多個上下文示例,我們提出了一種簡單但有效的上下文集成策略——特征集成(feature ensemble),該策略在多示例提示設置下能提升模型性能。此外,SegGPT 在不更新模型參數的前提下,也可以作為專用模型使用,通過調整特定的提示即可適配具體場景,例如在 ADE20K 數據集上進行語義分割。
我們的主要貢獻如下:
- 首次展示了一個通用模型,可以自動執行一系列多樣的分割任務;
- 無需微調,我們直接評估了預訓練的 SegGPT 在多個任務上的表現,包括小樣本語義分割、視頻目標分割、語義分割和全景分割;
- 實驗結果表明,SegGPT 在域內和跨域目標分割方面均具有出色的泛化能力,無論是定性還是定量。
需要指出的是,本研究的目標并不是在所有基準測試中超越已有的專用方法或達到新的 SOTA(state-of-the-art)性能,因為我們認為這并不是一個通用模型的責任所在。
2. 相關工作
2.1. 圖像分割
分割是計算機視覺中的一個基礎問題,其目標是在像素級別定位和組織有意義的概念。分割任務的類型根據概念的定義而異,例如foreground、category或object instance。例如,semantic segmentation [55] 涉及對圖像進行像素級的語義分類,而instance segmentation [30] 則旨在識別不同的object instance及其類別。video object segmentation [52, 37, 12] 的任務是:在給定視頻第一幀中目標mask的情況下,在整個視頻序列中對特定目標進行分割。
以往的分割方法 [32, 28, 54, 39, 18, 11, 2, 48, 23, 5, 8] 通常是針對特定任務設計的,不能靈活地切換任務或更換類別。本文提出了一個通用接口,該接口與所有分割任務兼容,通過恰當的訓練方案,一個單一的通用模型就能在in-domain和out-of-domain的分割任務上都獲得良好的性能,無論是從定性還是定量的角度看。
2.2. 通用視覺模型
近年來,人們嘗試使用基于Transformer的模型來統一視覺領域中的不同任務,從而誕生了一些通用視覺模型(vision generalists)[6, 7, 56, 33, 24]。DETR [5] 是最早使用Transformer [42] 作為目標檢測任務特定head的模型之一。Pix2Seq系列 [6, 7] 將視覺任務的輸出空間定義為離散的token序列,并以自回歸方式執行目標檢測、instance segmentation、關鍵點估計以及圖像描述等任務。Unified-IO [33] 和 OFA [45] 采用sequence-to-sequence方式對視覺、視覺-語言和NLP任務進行聯合建模,其中輸入和輸出都被定義為離散token的序列。UViM [24] 將panoptic segmentation、深度估計、圖像著色等像素標注任務統一到一起,但為每種任務單獨訓練模型。
盡管這些工作看起來是在將不同任務統一到相似的空間中,實際上它們仍通過某種形式的硬指示符(如特殊token)來完成每個任務,從而難以推廣到新任務。與之相對,本文提出的方法使用in-context框架,保持了任務定義的靈活性,并利用隨機上色方案來避免模型退化為多任務學習解決方案,而是通過引用上下文信息來完成指定任務。另一個區別在于任務范圍:本工作主要聚焦于視覺感知中的一個核心類別,即圖像分割。
2.3. In-Context視覺學習
GPT-3 [3] 將in-context learning引入深度學習,使得一系列NLP任務可以在給定prompt和示例的條件下被表述為文本補全問題。在計算機視覺領域,[1] 首次提出一種使用離散token進行圖像修復的in-context訓練框架,訓練數據來自視覺論文中的圖表和信息圖,展示了其在前景分割、單目標檢測和圖像著色方面的能力。Painter [46] 采用基于連續像素的masked image modeling方法,在七個多樣且具有挑戰性的視覺任務上進行in-context訓練,使用有監督數據集,并在這些任務上取得了競爭性的結果。
我們的工作構建在Painter框架之上,但專注于分割任務,因為其在視覺感知中扮演著核心角色。因此,本工作統一了多樣的分割數據,包括semantic segmentation、instance segmentation、part segmentation,甚至是特殊場景(如航空圖像)下的分割數據。此外,我們設計了一種隨機上色方案,迫使模型必須引用上下文信息才能完成任務,從而避免退化為多任務解法。由于分割任務和數據集的變化性小于深度/姿態估計,因此更容易共享內部結構以實現對in-domain任務的有效訓練,同時保持對out-of-domain分割任務的泛化能力。
3. 方法
SegGPT 是 Painter [46] 框架的一個特殊版本,它能夠通過一個通用的Painter對一切進行分割,因此命名為SegGPT。該訓練框架將視覺任務的輸出空間重新定義為“圖像”,并將不同任務統一為同一種圖像修復問題,即:隨機遮蓋任務輸出圖像并重建缺失像素。為了保持簡潔性與通用性,我們對架構與損失函數不作修改,即使用原始的ViT [13] 和簡單的smooth-?1 [17] 損失,但在in-context訓練中設計了一種新的隨機上色方案,以提升模型的泛化能力。
3.1. In-Context Coloring
在傳統的Painter框架中,每個任務的顏色空間是預定義的,這會導致模型的解法退化為多任務學習。例如,對于semantic segmentation,會預先定義一組顏色,并為每個語義類別分配固定的顏色。同樣地,在instance segmentation中,實例目標的顏色是根據其空間位置類別分配的,即顏色的數量等于空間位置的數量,這導致模型只依賴顏色本身來判斷任務,而不是利用不同segment之間的關系。
為了解決上述限制,我們提出了一種用于in-context coloring的隨機著色方案。我們首先隨機采樣另一張與輸入圖像具有相似上下文的圖像,例如具有相同的語義類別或object instance。接著,從目標圖像中隨機采樣一組顏色,并將每種顏色映射為一個隨機顏色,從而對對應像素進行重新著色。由此我們獲得兩對圖像,定義為一個in-context pair。
此外,我們引入了mix-context training方法,通過使用混合樣例對模型進行訓練。這一方法將多張具有相同顏色映射的圖像拼接在一起,然后對拼接后的圖像進行隨機裁剪和縮放,從而生成一個mixed-context訓練樣本。通過這種方式,模型學會關注圖像的上下文信息,而不是僅依賴特定的顏色信息來判斷任務。
這種統一方式使我們能夠以一致的方式利用所有segmentation數據集,僅根據具體任務調整數據采樣策略。我們根據不同的數據類型定義不同的上下文。例如,對于semantic segmentation,我們隨機采樣類別;對于instance segmentation,我們以隨機數量采樣object instances。同一圖像的不同視角,例如通過一組圖像增強方式變換而得的視角,也被視為上下文中的圖像。在實現中,采樣完全基于顏色,例如相同的顏色可表示相同的類別或相同的instance。
3.2. 上下文集成(Context Ensemble)
一旦訓練完成,SegGPT在推理過程中能夠釋放出其全部潛力。SegGPT支持在上下文中進行任意的分割任務,例如,使用單張圖像及其目標圖像的例子。目標圖像可以是單一顏色(不包括背景),也可以是多種顏色,例如在一次分割中分割多個類別或感興趣的對象。具體而言,給定一張待測試的輸入圖像,我們將其與示例圖像拼接,并將其輸入到SegGPT中,以獲得相應的in-context預測。
為了提供更準確和具體的上下文,可以使用多個示例。例如,可以使用相同語義類別的多個示例,或視頻中的前幾幀。為了有效地利用多個示例,我們提出了兩種上下文集成方法。一種是叫做空間集成(Spatial Ensemble),將多個示例按n × n網格拼接,然后重新采樣至與單個示例相同的大小。這種方法與in-context著色的直覺一致,多個示例的語義信息幾乎不增加額外成本就可以被提取到上下文中。另一種方法是特征集成(Feature Ensemble)。多個示例在批次維度上合并并獨立計算,唯一不同的是,在每個注意力層之后,查詢圖像的特征會被平均。通過這種方式,在推理過程中,查詢圖像可以從多個示例中收集信息。
3.3. 上下文調優(In-Context Tuning)
SegGPT能夠在不更新模型參數的情況下,適應特定的使用案例。我們凍結整個模型,并初始化一個可學習的圖像張量作為輸入上下文。在訓練過程中,只有這個可學習的圖像張量會被更新,其余部分的訓練保持不變,例如使用相同的損失函數。調優完成后,我們可以將學習到的圖像張量提取出來,并將其作為特定應用的即插即用鍵。例如,給定一個固定對象類別集的數據集,如ADE20K,我們可以為該數據集訓練一個定制的提示,而不會損害模型的通用性。或者,我們可以為特定場景(如你的公寓)或特定人物(如Bert的面部)優化一個提示圖像。這為廣泛的應用場景開辟了機會。
溫馨提示:
閱讀全文請訪問"AI深語解構" SegGPT: 在上下文中分割一切