fix(mlwm-train): train decoder on encoder residuals

Author: CEQ151

blind_watermark/mlwm/train.py

@@ -243,6 +243,7 @@ def train_main(args) -> dict[str, Any]:
     for stage_name, stage_cfg in iter_enabled_stages(config, args.stage):
       freeze_encoder = bool(stage_cfg.get('freeze_encoder', False))
       freeze_decoder = bool(stage_cfg.get('freeze_decoder', False))
+      use_fixed_residual = bool(stage_cfg.get('use_fixed_residual', freeze_encoder))
       stage_loss_cfg = deep_update(loss_cfg, stage_cfg.get('loss', {}))
       for parameter in encoder.parameters():
         parameter.requires_grad = not freeze_encoder
@@ -260,7 +261,7 @@ def train_main(args) -> dict[str, Any]:
           bits = batch['bits'].to(device)
 
           with torch.cuda.amp.autocast(enabled=scaler.is_enabled()):
-            if freeze_encoder:
+            if use_fixed_residual:
               residual = fixed_residual_batch(
                 torch,
                 bits,
@@ -377,14 +378,14 @@ def train_main(args) -> dict[str, Any]:
           for batch in val_loader:
             clean = batch['image'].to(device)
             bits = batch['bits'].to(device)
-            residual = encoder(clean, bits) if not freeze_encoder else fixed_residual_batch(
+            residual = fixed_residual_batch(
               torch,
               bits,
               clean.shape[-2],
               clean.shape[-1],
               float(stage_cfg.get('fixed_residual_scale', 0.02)),
               device,
-            )
+            ) if use_fixed_residual else encoder(clean, bits)
             watermarked = torch.clamp(clean + residual, 0.0, 1.0)
             attacked = attack_batch(torch, watermarked, attack_cfg, str(stage_cfg.get('attack_strength', 'medium'))).to(device)
             logits, confidence = decoder(attacked)

configs/mlwm/crop_decoder_5090.yaml

@@ -0,0 +1,46 @@
+extends: configs/mlwm/boost_5090.yaml
+optimizer:
+  lr: 0.00004
+  weight_decay: 0.00001
+stages:
+  warmup:
+    enabled: false
+  main:
+    enabled: false
+  hard_negative:
+    enabled: true
+    epochs: 18
+    freeze_encoder: true
+    use_fixed_residual: false
+    freeze_decoder: false
+    attack_strength: medium
+    loss:
+      payload_weight: 14.0
+      image_weight: 0.0
+      tv_weight: 0.0
+      residual_weight: 0.0
+  finalize:
+    enabled: true
+    epochs: 8
+    freeze_encoder: false
+    freeze_decoder: false
+    attack_strength: medium
+    loss:
+      payload_weight: 12.0
+      image_weight: 0.65
+      tv_weight: 0.018
+      residual_weight: 0.025
+attack:
+  enabled_ops:
+    - crop
+    - screenshot_sim
+    - color_jitter
+  ops_per_sample:
+    min: 1
+    max: 1
+  crop_keep: [0.96, 1.0]
+  resize_scale: [0.85, 1.15]
+  jpeg_quality: [75, 95]
+  webp_quality: [75, 95]
+  gaussian_blur_sigma: [0.0, 0.35]
+  gaussian_noise_std: [0.0, 0.004]