1. 從“輸入和標簽都是圖片”說起去雨任務的數(shù)據(jù)集到底特殊在哪圖像去雨這個方向很多人第一次接觸時都會卡在一個很樸素的問題上數(shù)據(jù)集到底長什么樣如果你之前做過分類任務習慣的是“一張圖 一個類別標簽”做過分割任務習慣的是“一張圖 一張掩碼圖”。而去雨任務屬于圖像到圖像的映射輸入是一張帶雨痕的圖標簽是同一場景下沒有雨的那張干凈圖。聽起來簡單但真正動手搭 Dataset 的時候坑遠比想象中多。我最初拿到一個去雨數(shù)據(jù)集時第一反應是“不就是兩個文件夾一個放帶雨圖一個放干凈圖按文件名配對嘛”。結果打開目錄一看帶雨圖有 1200 張干凈圖只有 400 張文件名還完全對不上。這就是去雨數(shù)據(jù)集最典型的特征之一它不是嚴格的一一對應關系而是“一張干凈圖對應多張不同雨痕的合成圖”。這個結構直接決定了你在寫 PyTorch 的Dataset類時配對邏輯不能想當然。為什么會出現(xiàn)這種結構因為真實世界里“同一場景、同一角度、同一光照只有雨不一樣”的成對數(shù)據(jù)幾乎拍不到。雨是動態(tài)的光線在變相機稍微抖一下場景就變了。所以主流做法是先采集干凈背景圖然后用物理渲染或者圖層疊加的方式人工合成不同密度、不同方向、不同雨滴大小的雨痕生成多張帶雨圖。這樣一來一張干凈圖就派生出了好幾張訓練樣本。理解這一點非常關鍵因為它影響三件事第一你的數(shù)據(jù)配對邏輯必須支持“一對多”第二劃分訓練集和驗證集時必須按“干凈圖”來劃分而不是按“帶雨圖”劃分否則同一張干凈圖派生出的帶雨圖會同時出現(xiàn)在訓練和驗證里造成數(shù)據(jù)泄漏第三評估指標的可信度依賴于驗證集里干凈圖的多樣性而不是帶雨圖的數(shù)量。我見過不少新手在這里翻車驗證集 loss 低得離譜以為模型學得很好結果一放到真實雨天照片上就完全失效。排查半天才發(fā)現(xiàn)訓練集和驗證集里存在同一場景的不同雨痕版本模型其實是在“背場景”而不是在“學去雨”。這個坑非常隱蔽因為從帶雨圖的文件名上根本看不出來它們共享同一張干凈背景。所以這一節(jié)我想先把認知層面的事情講透去雨數(shù)據(jù)集的核心不是“圖片數(shù)量”而是“場景數(shù)量”和“雨痕多樣性”的平衡。你后面寫 Dataset、做增強、劃分數(shù)據(jù)集所有決策都要圍繞這個認知展開。下面我會從目錄結構設計、配對邏輯、數(shù)據(jù)劃分三個角度把這件事拆開講。1.1 常見的兩種目錄組織方式與各自的坑去雨數(shù)據(jù)集在磁盤上的組織方式我見過的主流有兩種。第一種是雙文件夾平行結構dataset/ rainy/ scene01_rain_a.png scene01_rain_b.png scene02_rain_a.png ... clean/ scene01.png scene02.png ...這種結構的特點是干凈圖只有一份帶雨圖有多份通過文件名前綴比如scene01來關聯(lián)。寫 Dataset 時你需要從帶雨圖文件名里解析出場景 ID再去 clean 目錄里找對應的干凈圖。優(yōu)點是磁盤占用小干凈圖不重復存儲缺點是配對邏輯依賴命名規(guī)范一旦命名不統(tǒng)一就會出錯。第二種是成對文件夾結構dataset/ train/ rainy/ scene01_rain_a.png scene01_rain_b.png clean/ scene01_rain_a.png scene01_rain_b.png注意這里 clean 目錄里其實是同一張圖復制了兩份只是文件名和帶雨圖一一對應。這種結構寫 Dataset 最簡單直接按文件名配對即可但磁盤占用會成倍增加。如果你的干凈圖是 2K 分辨率、幾千張復制一份可能就多占幾個 G。我個人的選擇是訓練階段用第一種配合內(nèi)存緩存如果磁盤實在緊張第二種只在小規(guī)模調試時用。原因很簡單去雨模型的訓練輪次通常很多磁盤 IO 會成為瓶頸與其復制文件不如在 Dataset 里做一次映射緩存。這里有個細節(jié)要注意解析場景 ID 時不要用簡單的字符串切割。比如scene01_rain_a.png和scene01_rain_b.png如果你用split(_)[0]拿到scene01看起來沒問題但如果有的數(shù)據(jù)集命名是scene_01_rain.png切割結果就錯了。穩(wěn)妥的做法是用正則表達式提取“場景標識 雨痕標識”兩部分并且把映射關系在初始化時一次性建好存成一個字典避免每個__getitem__都去遍歷目錄。1.2 數(shù)據(jù)劃分為什么必須按干凈圖切分前面提到按干凈圖劃分這里展開說清楚操作方式。假設你有 400 張干凈圖每張派生出 3 張帶雨圖總共 1200 張帶雨圖。正確的劃分方式是先把 400 張干凈圖按 8:1:1 切成訓練 320、驗證 40、測試 40然后每張干凈圖對應的 3 張帶雨圖跟著它一起進對應的集合。這樣訓練集有 960 張帶雨圖驗證集和測試集各 120 張。錯誤的方式是直接把 1200 張帶雨圖打亂按比例切分。這樣同一張干凈圖的 3 張帶雨圖很可能被分到不同集合驗證集里的場景在訓練集里出現(xiàn)過指標虛高。實現(xiàn)上我通常會在 Dataset 初始化之前先構建一個scene_id - [rainy_paths]的字典然后對scene_id列表做劃分最后展開成實際的圖片路徑列表。這樣邏輯清晰也不容易出錯。劃分時記得固定隨機種子保證每次實驗的劃分一致否則不同輪次的對比就沒有意義了。提示如果你的數(shù)據(jù)集本身沒有提供場景 ID只能從文件名推斷那一定要先寫個小腳本統(tǒng)計一下“干凈圖數(shù)量”和“帶雨圖數(shù)量”的比例以及文件名前綴的分布情況。我遇到過前綴解析錯誤導致 1200 張圖只匹配上 300 張的情況訓練時 loss 正常下降但模型效果極差排查了很久才發(fā)現(xiàn)是配對漏了。2. 寫一個能跑通的去雨 Dataset從__init__到__getitem__認知理清之后進入代碼層面。PyTorch 的Dataset類看起來簡單就三個方法但去雨任務的特殊性讓每個方法都有講究。我下面給出的實現(xiàn)是我自己在多個項目里迭代過的版本兼顧了靈活性、性能和可讀性。先明確目標輸入是帶雨圖路徑輸出是(rainy_tensor, clean_tensor)這樣一對張量形狀都是[3, H, W]數(shù)值范圍歸一化到[0, 1]或者標準化到[-1, 1]具體取決于你后面用的損失函數(shù)和預訓練模型。2.1__init__里該做什么、不該做什么__init__的職責是建立索引和映射而不是加載圖片數(shù)據(jù)。這一點很多人會搞混把圖片讀取放在__init__里結果數(shù)據(jù)集一大內(nèi)存直接爆掉或者啟動時間長得離譜。我的做法是在__init__里完成這幾件事第一掃描目錄構建rainy_paths和clean_paths兩個列表并且保證它們一一對應。如果是“一對多”結構rainy_paths會比clean_paths長但每個rainy_path都能通過映射找到對應的clean_path。第二解析場景 ID建立scene_to_indices字典方便后續(xù)按場景劃分。第三接收transform參數(shù)。注意去雨任務的 transform 有個特殊要求輸入和標簽必須做完全相同的幾何變換。如果你用torchvision.transforms里的RandomCrop分別作用在兩張圖上很可能裁到不同位置配對就廢了。正確做法是用torchvision.transforms.functional或者自己寫一個同時接收兩張圖的 transform 函數(shù)。第四可選地接收一個cache參數(shù)。如果數(shù)據(jù)集不大比如干凈圖只有幾百張可以在__init__里把干凈圖預加載到內(nèi)存帶雨圖按需讀取。這樣能顯著減少 IO因為干凈圖會被反復讀取。這里有個經(jīng)驗不要用os.listdir直接拿文件列表因為它不保證順序。如果你先listdir帶雨圖目錄再listdir干凈圖目錄兩個列表的順序可能不一致導致配對錯亂。穩(wěn)妥的做法是用sorted(os.listdir(...))或者自己寫一個按場景 ID 排序的邏輯。2.2__getitem__的配對邏輯與異常處理__getitem__是核心它接收一個索引返回一對張量。我的實現(xiàn)大致是這樣的思路def __getitem__(self, idx): rainy_path self.rainy_paths[idx] scene_id self.parse_scene_id(rainy_path) clean_path self.scene_to_clean[scene_id] rainy_img Image.open(rainy_path).convert(RGB) clean_img Image.open(clean_path).convert(RGB) # 同步幾何變換 rainy_tensor, clean_tensor self.joint_transform(rainy_img, clean_img) return rainy_tensor, clean_tensor看起來簡單但有幾個坑必須處理?？右粓D片尺寸不一致。有些數(shù)據(jù)集的帶雨圖和干凈圖分辨率不同比如帶雨圖是裁剪過的干凈圖是原圖。直接送進網(wǎng)絡會報錯。我的做法是在__init__里做一次尺寸檢查如果發(fā)現(xiàn)不一致要么統(tǒng)一 resize要么在__getitem__里做中心裁剪。中心裁剪更安全因為它不會引入插值偽影而去雨任務對偽影很敏感?？佣D片模式不是 RGB。有些干凈圖是灰度圖或者帶 alpha 通道的 PNG。convert(RGB)能解決大部分問題但如果原圖是 16 位深度convert會截斷成 8 位損失信息。如果你的數(shù)據(jù)集是 16 位需要用numpy讀取并手動歸一化?？尤募p壞或缺失。大規(guī)模數(shù)據(jù)集里偶爾會有損壞的圖片。如果__getitem__直接拋異常訓練會中斷。我的做法是用try-except包住讀取邏輯出錯時隨機返回另一對樣本并打印警告。這樣訓練不會斷但你要定期檢查警告日志看看是不是有系統(tǒng)性問題?？铀呐鋵κ　Ｈ绻硞€帶雨圖的場景 ID 在干凈圖里找不到說明數(shù)據(jù)集本身有問題。這種情況我建議在__init__階段就檢查并報錯而不是等到訓練時才發(fā)現(xiàn)?？梢栽赺_init__里遍歷所有帶雨圖驗證每個場景 ID 都有對應的干凈圖缺失的直接列出來。2.3 同步幾何變換去雨任務最容易寫錯的地方這一小節(jié)單獨拿出來講因為它太重要了。去雨任務的輸入和標簽是同一場景的兩張圖空間位置必須嚴格對齊。如果你對輸入做隨機裁剪對標簽做另一套隨機裁剪兩張圖就錯位了模型學到的是錯誤映射loss 可能下降但效果完全不對。正確的做法是生成一組隨機參數(shù)然后把這組參數(shù)同時應用到兩張圖上。比如隨機裁剪你先隨機生成裁剪的左上角坐標和寬高然后用同一組坐標去裁兩張圖。隨機翻轉同理生成一個布爾值兩張圖一起翻或不翻。torchvision.transforms.functional提供了crop、hflip、vflip、rotate等函數(shù)它們接收具體的參數(shù)而不是隨機生成參數(shù)正好適合這種場景。你可以自己封裝一個JointTransform類在__call__里生成隨機參數(shù)并應用到兩張圖。如果你用的是albumentations庫它原生支持additional_targets可以把標簽圖作為額外目標保證變換同步。配置方式是在Compose里加additional_targets{image0: image}然后調用時傳入imagerainy, image0clean。這個庫在處理同步變換上比手寫更省心但要注意版本兼容性不同版本的 API 有差異。注意不要用RandomResizedCrop這類會改變圖像內(nèi)容的變換去做同步因為兩張圖如果內(nèi)容不同一張有雨一張沒雨resize 后的插值結果會有細微差異雖然影響不大但嚴格來說不如直接裁剪干凈。我通常只用隨機裁剪 隨機翻轉最多加一個 90 度旋轉。3. 數(shù)據(jù)加載性能為什么你的 GPU 利用率上不去模型訓練慢很多人第一反應是模型太大或者 GPU 不行。但在去雨任務里更常見的原因是數(shù)據(jù)加載成了瓶頸。去雨模型的輸入輸出都是高分辨率圖片讀取、解碼、變換的開銷很大如果 DataLoader 配置不當GPU 大部分時間在等數(shù)據(jù)。我做過一個對比實驗同樣的模型、同樣的數(shù)據(jù)量只調整 DataLoader 的參數(shù)訓練速度差了將近一倍。下面把我踩過的坑和調優(yōu)經(jīng)驗整理出來。3.1num_workers不是越大越好num_workers控制數(shù)據(jù)加載的并行進程數(shù)。理論上越大越快但實際上受限于 CPU 核心數(shù)、磁盤 IO 和內(nèi)存帶寬。我試過在 8 核機器上設num_workers16結果反而比num_workers4慢因為進程切換和內(nèi)存拷貝的開銷超過了并行收益。我的經(jīng)驗值是num_workers設為 CPU 物理核心數(shù)的 1 到 1.5 倍。比如 8 核就設 8 或 12。如果是 SSD可以適當調大如果是機械硬盤調大反而會因為隨機讀取導致尋道開銷增加。另外num_workers 0時Dataset 里的對象會被復制到子進程如果你在__init__里緩存了大量數(shù)據(jù)每個 worker 都會復制一份內(nèi)存占用會成倍增加。解決辦法是用torch.utils.data.DataLoader的persistent_workersTrue配合prefetch_factor或者把緩存改成共享內(nèi)存。3.2pin_memory和prefetch_factor的配合pin_memoryTrue會把數(shù)據(jù)放到鎖頁內(nèi)存加速 CPU 到 GPU 的傳輸。這個參數(shù)在 GPU 訓練時幾乎必開但要注意它和num_workers的配合。如果num_workers0pin_memory的效果有限因為數(shù)據(jù)加載和傳輸是串行的。prefetch_factor控制每個 worker 預取多少個 batch。默認是 2對于去雨這種 IO 密集的任務可以調到 4 或 6。但調太大也會占內(nèi)存需要根據(jù)你的圖片分辨率和 batch size 權衡。我通常的配置是loader DataLoader( dataset, batch_size8, shuffleTrue, num_workers8, pin_memoryTrue, prefetch_factor4, persistent_workersTrue, drop_lastTrue )drop_lastTrue在去雨任務里建議開啟因為最后一個 batch 如果只有一張圖BatchNorm 的統(tǒng)計會不穩(wěn)定影響訓練。3.3 圖片解碼的優(yōu)化從 PIL 到 OpenCV 再到 DALIPIL 是 PyTorch 生態(tài)里最常用的圖片讀取庫但它的解碼速度一般。如果你的數(shù)據(jù)集很大可以換成 OpenCV 的cv2.imread速度通?？?2 到 3 倍。注意 OpenCV 讀進來是 BGR 順序需要轉成 RGB。更進一步如果 CPU 解碼仍然是瓶頸可以考慮 NVIDIA 的 DALI 庫它能把解碼放到 GPU 上做。但 DALI 的配置比較復雜而且對自定義變換的支持不如 PIL 靈活。我的建議是先用 PIL 跑通確認瓶頸在數(shù)據(jù)加載后再考慮換庫。很多時候把num_workers和prefetch_factor調好問題就解決了。還有一個容易被忽略的點圖片格式。PNG 是無損壓縮解碼慢但質量高JPEG 解碼快但有壓縮偽影。去雨任務對偽影敏感我建議訓練時用 PNG如果磁盤空間不夠可以用高質量的 JPEGquality 95 以上。我實測過quality 95 的 JPEG 和 PNG 在去雨模型上的效果差異很小但讀取速度快不少。4. 訓練中的那些坑從 loss 不下降到雨痕殘留數(shù)據(jù)和加載都搞定之后訓練過程本身也有一堆坑。去雨任務和普通圖像恢復任務相比有幾個特殊之處雨痕是疊加在圖像上的不是簡單遮擋雨痕有方向性不同方向的雨需要不同的處理雨痕的密度變化很大從稀疏雨滴到暴雨都有。這些特性決定了你的訓練策略不能照搬超分或者去噪。4.1 損失函數(shù)的選擇L1、L2 還是感知損失L2 損失MSE在去雨任務里有個著名的問題它傾向于產(chǎn)生模糊的結果。因為 L2 對誤差的懲罰是平方的模型會傾向于預測所有可能值的平均導致雨痕區(qū)域被“抹平”而不是“去除”。你去看用純 L2 訓練的去雨模型雨痕是沒了但原本的紋理也糊了。L1 損失比 L2 好一些因為它對異常值的懲罰是線性的不會過度平滑。但 L1 也有問題它在零點不可導訓練后期容易震蕩。實踐中我通常用SmoothL1Loss也叫 Huber Loss它在誤差小的時候表現(xiàn)為 L2誤差大的時候表現(xiàn)為 L1兼顧了穩(wěn)定性和銳度。如果你追求更好的視覺效果可以加感知損失Perceptual Loss用預訓練的 VGG 提取特征比較去雨結果和干凈圖在特征空間的差異。感知損失能讓結果更符合人眼觀感但訓練成本會上升而且需要小心權重設置感知損失權重太大時PSNR 會下降。我的經(jīng)驗配置是L1 0.1 * Perceptual感知損失用 VGG16 的 relu2_2 和 relu4_3 層。這個配置在多個數(shù)據(jù)集上表現(xiàn)穩(wěn)定。4.2 雨痕殘留和過度去雨的平衡訓練去雨模型時你會遇到兩種相反的問題一種是雨痕沒去干凈另一種是把原本的紋理當成雨去掉了導致圖像過度平滑。雨痕殘留通常發(fā)生在訓練早期模型還沒學會識別雨痕。這時候可以檢查一下學習率是不是太大或者數(shù)據(jù)增強是不是太激進。如果雨痕是細密的斜線而你的隨機裁剪把大部分雨痕裁掉了模型就學不到。解決辦法是確保裁剪后的區(qū)域仍然包含足夠的雨痕可以通過統(tǒng)計雨痕像素比例來篩選。過度去雨通常發(fā)生在訓練后期模型為了降低 loss把一些高頻紋理也當成雨去掉了。這時候可以降低學習率或者引入一個“保真度”約束比如在干凈圖的高頻區(qū)域加大 loss 權重。另一個技巧是使用對抗損失GAN Loss讓判別器判斷結果是否像真實干凈圖但這會顯著增加訓練復雜度新手不建議一上來就用。4.3 驗證集指標和肉眼觀感的差距PSNR 和 SSIM 是去雨任務最常用的指標但它們和肉眼觀感經(jīng)常不一致。我遇到過 PSNR 很高的模型結果圖看起來還是有霧感也遇到過 PSNR 一般但觀感很干凈的模型。原因是 PSNR 對亮度偏移很敏感如果模型整體偏亮或偏暗PSNR 會掉但人眼可能覺得還好。SSIM 好一些但它對紋理的敏感度也不夠。我的做法是訓練時看 PSNR/SSIM但最終選模型一定要肉眼對比。我會在驗證集里挑幾張有代表性的圖稀疏雨、密集雨、不同背景把不同 epoch 的結果拼在一起看。有時候 PSNR 最高的 epoch 并不是觀感最好的選那個觀感好的指標差 0.1 到 0.2 完全可以接受。另外驗證時要注意把結果保存成圖片不要只看數(shù)值。我習慣每幾個 epoch 就導出一次驗證結果用圖片瀏覽器快速翻看比盯著數(shù)字直觀得多。5. 從單張推理到批量部署實際使用中的工程問題模型訓練好之后真正的挑戰(zhàn)才開始。研究階段你處理的是固定尺寸的圖片實際使用中你會遇到各種尺寸、各種來源的圖片。這一節(jié)講講工程化落地時的問題。5.1 任意尺寸輸入的推理策略去雨模型通常是全卷積網(wǎng)絡理論上支持任意尺寸輸入。但實際推理時如果圖片太大顯存會爆如果圖片太小感受野不夠效果會下降。我的策略是把大圖切成有重疊的 patch分別推理后拼接。重疊區(qū)域取平均避免拼接縫。patch 大小根據(jù)你的模型和顯存來定我通常用 512x512重疊 64 像素。這個配置在 8G 顯存的卡上跑得很穩(wěn)。切 patch 時要注意不要簡單地從左上角開始切因為雨痕可能正好落在切割線上。有重疊的切法能緩解這個問題但也不能完全避免。如果對質量要求極高可以用滑動窗口步長小于 patch 大小重疊區(qū)域用高斯權重融合。5.2 批量推理時的顯存管理批量推理時顯存占用和 batch size 成正比。如果你要處理幾千張圖不要一次性全部加載而是用 DataLoader 流式讀取推理完一張釋放一張。另外推理時記得用torch.no_grad()否則會保存計算圖顯存占用會大很多。如果模型里有 BatchNorm推理時要設model.eval()否則會使用當前 batch 的統(tǒng)計量結果不穩(wěn)定。我見過有人推理時忘了eval()結果單張推理和批量推理的結果不一致排查了半天。這個坑很隱蔽因為訓練時一切正常只有推理時才暴露。5.3 結果后處理銳化和色彩校正模型輸出的圖片有時候會偏軟或者偏色。偏軟是因為 L1/L2 損失的平滑效應偏色是因為訓練數(shù)據(jù)的色彩分布和實際圖片不一致。后處理可以緩解這些問題。銳化用簡單的UnsharpMask就行強度不要太大否則會放大噪聲。色彩校正可以用直方圖匹配把結果的色彩分布對齊到輸入圖注意是對齊到輸入圖不是干凈圖因為實際使用時你沒有干凈圖。但后處理要謹慎它可能掩蓋模型本身的問題。我的建議是先看原始輸出如果問題不嚴重不要加后處理如果確實需要把后處理參數(shù)記錄下來方便復現(xiàn)。6. 一些零散但重要的經(jīng)驗最后分享幾個我在做去雨項目時積累的零散經(jīng)驗它們不屬于某個特定環(huán)節(jié)但每個都幫我省過時間。關于數(shù)據(jù)增強去雨任務的增強要克制。隨機裁剪和翻轉是安全的但顏色抖動要小心因為顏色變化可能讓模型把顏色異常誤認為雨痕。我通常只做輕微的亮度調整幅度控制在 0.1 以內(nèi)。關于預訓練模型如果你的去雨網(wǎng)絡用了預訓練 backbone比如 ResNet加載權重時要注意有些預訓練模型是在 ImageNet 上訓練的輸入歸一化參數(shù)不同。用錯歸一化參數(shù)會導致訓練初期 loss 爆炸。關于隨機種子去雨任務的隨機性來源很多數(shù)據(jù)劃分、數(shù)據(jù)增強、權重初始化固定種子能讓你復現(xiàn)實驗。但要注意num_workers 0時DataLoader 的隨機種子需要在 worker 初始化函數(shù)里設置否則每個 worker 的隨機序列可能相同導致增強結果重復。關于日志訓練去雨模型時建議把每個 epoch 的驗證結果圖保存下來按 epoch 編號命名。這樣訓練結束后可以快速回看模型是怎么逐步變好的也能發(fā)現(xiàn)某些 epoch 的異常比如突然變差又恢復可能是學習率調度的問題。關于顯存去雨模型的顯存占用和圖片分辨率是平方關系。512x512 的圖batch size 8 可能占 6G 顯存1024x1024 的圖同樣的 batch size 可能就要 20G 以上。如果你的顯存不夠優(yōu)先降 batch size而不是降分辨率因為分辨率對去雨效果的影響更大。關于評估不要只在合成數(shù)據(jù)集上評估。合成雨和真實雨在物理特性上有差距合成雨通常是均勻的斜線真實雨有雨滴、有霧、有鏡頭上的水漬。如果你的模型只在合成數(shù)據(jù)上訓練放到真實照片上效果會打折扣。有條件的話收集一些真實雨天照片做定性評估哪怕沒有干凈圖對照肉眼看看雨痕去得干不干凈也是有價值的。我在實際項目里最深的一個體會是去雨任務的上限往往不在模型結構而在數(shù)據(jù)質量。一個簡單的 U-Net配上高質量、多樣化的成對數(shù)據(jù)效果可能比一個復雜網(wǎng)絡配劣質數(shù)據(jù)好得多。所以如果你剛開始做去雨先把數(shù)據(jù)集的配對、劃分、加載這三件事做扎實比急著換模型架構重要得多。