Scorer as Optimization Target

src/camera-based-e2e/models/base_model.py

@@ -3,6 +3,7 @@
 import torch.nn.functional as F
 import pytorch_lightning as pl
 import torchvision
+from contextlib import nullcontext
 from dataclasses import asdict, is_dataclass
 
 from .losses.depth_loss import DepthLoss
@@ -26,6 +27,9 @@ def __init__(self, model: nn.Module, lr: float, lr_vision: float | None = None,
         super(LitModel, self).__init__()
         self.model = model
 
+        # NVJPEG fall back if we are running on Negishi AMD GPU
+        self.has_nvjpeg = True
+
         # If we are using ScorerModel, which has a cfg, then save the attributes of the cfg as hparams, so they go into wandb
         cfg = getattr(model, "cfg", None)
         if cfg is None:
@@ -81,31 +85,47 @@ def decode_batch_jpeg(
         self,
         images_jpeg: list[list[torch.Tensor]],
         device: torch.device | None = None,
-    ) -> list[torch.Tensor]:
+    ) -> list[torch.Tensor | None]:
+        cam_idxs_used = tuple(getattr(self.model.cfg, "cam_idxs_used", range(len(images_jpeg))))
         decode_device = self.device if device is None else device
+
+        selected = [(cam_idx, images_jpeg[cam_idx]) for cam_idx in cam_idxs_used]
+
         # Flatten cameras
         flat_encoded, cam_sizes = [], []
-        for cam in images_jpeg:
-            cam_sizes.append(len(cam))
+        for cam_idx, cam in selected:
+            cam_sizes.append((cam_idx, len(cam)))
             for jpg in cam:
                 t = jpg if isinstance(jpg, torch.Tensor) else torch.frombuffer(memoryview(jpg), dtype=torch.uint8)
                 # decode_jpeg requires the raw jpeg bytes to be on cpu
                 if t.device.type != "cpu":
                     t = t.cpu()
                 flat_encoded.append(t)
 
-        flat_decoded = torchvision.io.decode_jpeg(
-            flat_encoded, 
-            mode=torchvision.io.ImageReadMode.UNCHANGED,
-            device=decode_device,
-        ) # list of (C, H, W) gpu tensors
+        try:
+            flat_decoded = torchvision.io.decode_jpeg(
+                flat_encoded, 
+                mode=torchvision.io.ImageReadMode.UNCHANGED,
+                device=decode_device,
+            ) # list of (C, H, W) gpu tensors
+        except Exception as e:
+            if "nvJPEG" not in str(e):
+                raise e
+            self.has_nvjpeg = False
+            flat_decoded = torchvision.io.decode_jpeg(
+                flat_encoded,
+                mode=torchvision.io.ImageReadMode.UNCHANGED,
+                device="cpu",
+            )
+            if torch.device(decode_device).type == "cuda":
+                flat_decoded = [img.to(decode_device, non_blocking=True) for img in flat_decoded]
 
-        out = []
+        out = [None] * len(images_jpeg)
         idx = 0
-        for n in cam_sizes:
-            cam_list = flat_decoded[idx: idx+n]
+        for cam_idx, n in cam_sizes:
+            out[cam_idx] = torch.stack(flat_decoded[idx:idx+n], dim=0)
             idx += n
-            out.append(torch.stack(cam_list, dim=0))  # (B, C, H, W)
+
         return out
 
     def on_fit_start(self) -> None:
@@ -211,10 +231,10 @@ def configure_optimizers(self):
         return optimizer
 
     # ---- forward / step ----
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
         return self.model(x)
 
-    def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor:
+    def _shared_step(self, batch: dict[str, torch.Tensor | list[torch.Tensor]], stage: str) -> torch.Tensor:
         past, future, intent = batch['PAST'], batch['FUTURE'], batch['INTENT']
 
         if "IMAGES" in batch:
@@ -234,8 +254,15 @@ def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor:
         pred_future = self.forward(model_inputs)  # (B, T*2)
         pred_depth = None
         pred_scores: torch.Tensor = None
+        pred_traj_flat: torch.Tensor = None
+        query_for_score: torch.Tensor = None
         if isinstance(pred_future, dict):
-            pred_future, pred_depth, pred_scores = pred_future["trajectory"], pred_future.get("depth", None), pred_future.get("scores", None)
+            outputs = pred_future
+            pred_future = outputs["trajectory"]
+            pred_depth = outputs.get("depth", None)
+            pred_scores = outputs.get("scores", None)
+            pred_traj_flat = outputs.get("trajectory_flat", None)
+            query_for_score = outputs.get("query_for_score", None)
 
         pred = pred_future
         t_steps = future.shape[1]
@@ -310,6 +337,67 @@ def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor:
         else:
             loss_score = torch.tensor(0.0, device=self.device)
 
+        # Train-only adversarial scorer supervision in trajectory space.
+        loss_adv = torch.tensor(0.0, device=self.device)
+        adv_enabled = bool(getattr(self.hparams, "model_cfg_adv_enabled", False))
+        if (stage == "train" and adv_enabled and k_modes > 1 and pred_scores is not None and pred_traj_flat is not None and query_for_score is not None and hasattr(self.model, "score_trajectories")):
+            epsilon = float(getattr(self.hparams, "model_cfg_adv_epsilon", 0.10)) # max perturbation
+            adv_steps = max(1, int(getattr(self.hparams, "model_cfg_adv_steps", 3))) # running Projected Gradient Descent
+            alpha = epsilon / adv_steps
+
+            original_traj = pred_traj_flat.detach()
+            adv_traj = original_traj
+            for _ in range(adv_steps):
+                # run forward pass of scorer on adversarial trajectory
+                attack_traj = adv_traj.detach().requires_grad_(True)
+                score_adv = self.model.score_trajectories(attack_traj, query_for_score)
+
+                # compute gradient, gradient ASCENT to MAXIMIZE loss.
+                grad = torch.autograd.grad(score_adv.sum(), attack_traj, only_inputs=True)[0]
+                adv_traj = attack_traj + alpha * grad.sign()
+                delta = torch.clamp(adv_traj - original_traj, min=-epsilon, max=epsilon)
+                adv_traj = (original_traj + delta).detach()
+
+            bsz, _, t2_adv = adv_traj.shape
+            t_adv = t2_adv // 2
+            adv_ade = torch.norm(adv_traj.view(bsz, k_modes, t_adv, 2) - future[:, None], dim=-1,).mean(dim=-1).detach()
+            adv_scores = self.model.score_trajectories(adv_traj, query_for_score)
+            loss_adv = F.mse_loss(adv_scores, adv_ade)
+
+        # Sobolev training. d(Score)/d(traj) ~= d(ADE)/d(traj)
+        loss_sobolev = torch.tensor(0.0, device=self.device)
+        sobolev_grad_cosine = torch.tensor(0.0, device=self.device)
+        sobolev_enabled = bool(getattr(self.hparams, "model_cfg_sobolev_enabled", False))
+        if (stage == "train" and sobolev_enabled and k_modes > 1 and pred_scores is not None and pred_traj_flat is not None and query_for_score is not None and hasattr(self.model, "score_trajectories")):
+            autocast_device = pred_traj_flat.device.type
+            autocast_ctx = (
+                torch.autocast(device_type=autocast_device, enabled=False)
+                if autocast_device in ("cpu", "cuda")
+                else nullcontext()
+            )
+            with autocast_ctx:
+                sobolev_traj = pred_traj_flat.detach().float().requires_grad_(True)
+                sobolev_query = query_for_score.detach().float()
+                sobolev_scores = self.model.score_trajectories(sobolev_traj, sobolev_query)
+                grad_score = torch.autograd.grad(
+                    sobolev_scores.sum(),
+                    sobolev_traj,
+                    create_graph=True,
+                    only_inputs=True,
+                )[0]
+
+                sobolev_traj_xy = sobolev_traj.view(pred.size(0), k_modes, t_steps, 2)
+                delta = sobolev_traj_xy - future[:, None].float()
+                dist = torch.norm(delta, dim=-1, keepdim=True).clamp_min(1e-3)
+                grad_ade = (delta / (t_steps * dist)).reshape_as(sobolev_traj)
+
+                loss_sobolev = F.smooth_l1_loss(grad_score, grad_ade.detach())
+                sobolev_grad_cosine = F.cosine_similarity(
+                    grad_score.reshape(pred.size(0), k_modes, -1),
+                    grad_ade.reshape(pred.size(0), k_modes, -1),
+                    dim=-1,
+                ).mean()
+
         # Scorer Metrics
         scorer_metrics = {}
         if k_modes > 1 and pred_scores is not None:
@@ -331,7 +419,8 @@ def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor:
             scorer_metrics[f"{stage}_scorer_spearman"] = rho.mean()
 
         # Depth Loss
-        if pred_depth is not None:
+        use_depth_loss = bool(getattr(self.hparams, "model_cfg_use_depth_loss", False))
+        if pred_depth is not None and use_depth_loss:
             front_img = images[1]  # front camera
             depth_in = F.interpolate(front_img, size=(128, 128), mode='nearest')
             loss_depth = self.depth_loss(depth_in, pred_depth, loss_fn=F.l1_loss)
@@ -341,13 +430,29 @@ def _shared_step(self, batch: torch.Tensor, stage: str) -> torch.Tensor:
         loss_depth *= 0.1 # slightly enabled
         loss_ade *= 1.0 # TODO: tune loss terms
         loss_score *= 1.0
-        total_loss = loss_ade + loss_depth + loss_score + loss_rfs
+        adv_lambda = float(getattr(self.hparams, "model_cfg_adv_lambda", 0.1))
+        sobolev_lambda = float(getattr(self.hparams, "model_cfg_sobolev_lambda", 0.1))
+        local_step_lambda = float(getattr(self.hparams, "model_cfg_local_step_lambda", 0.1))
+        total_loss = (
+            loss_ade
+            + loss_depth
+            + loss_score
+            + loss_rfs
+            + (adv_lambda * loss_adv)
+            + (sobolev_lambda * loss_sobolev)
+            + (local_step_lambda * loss_local_step)
+        )
         # TODO: improve logging both to disk and to console
         log_payload = {
             f"{stage}_loss_ade": loss_ade,
             f"{stage}_loss_score": loss_score,
             f"{stage}_loss_depth": loss_depth,
             f"{stage}_loss_rfs": loss_rfs,
+            f"{stage}_loss_adv": loss_adv,
+            f"{stage}_loss_sobolev": loss_sobolev,
+            f"{stage}_sobolev_grad_cosine": sobolev_grad_cosine,
+            f"{stage}_loss_local_step": loss_local_step,
+            f"{stage}_local_step_rel_improve": local_step_rel_improve,
             f"{stage}_rfs_unweighted": rfs_unweighted,
             f"{stage}_loss": total_loss,
         }

src/camera-based-e2e/models/feature_extractors.py

@@ -2,6 +2,8 @@
 import torch
 import torch.nn as nn
 import timm
+import torchvision
+from torchvision.transforms import v2
 
 class DINOFeatures(nn.Module):
     def __init__(self, model_name: str = "vit_small_plus_patch16_dinov3.lvd1689m", frozen: bool = True):
@@ -14,15 +16,15 @@ def __init__(self, model_name: str = "vit_small_plus_patch16_dinov3.lvd1689m", f
             for param in self.dino_model.parameters():
                 param.requires_grad = False
 
-        self.dims = [384, 384, 384]  # feature dims for each layer
+        self.dims = [384]  # feature dims for last layer
         self.patch_size = 16  # patch size
 
     def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
         # x: (B, 3, H, W)
         # transforms: resize 256x256, center crop, normalize
         x_t = self.transforms(x.float().div(255.0)) # preprocess
         features = self.dino_model(x_t)
-        return features # 3 x [B, 384, 16, 16]
+        return features[-1:] # return last layer features as list of 1 tensor (B, C, H', W')
 
 class SAMFeatures(nn.Module):
     def __init__(
@@ -52,4 +54,48 @@ def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
         # x: (B, 3, H, W)
         x_t = self.transforms(x.float().div(255.0))  # preprocess
         feats = self.sam_model(x_t)       # list of feature maps
-        return [feats[self.feature_stage]]
+        return [feats[self.feature_stage]] # (B, C, H', W')
+
+class EUPEFeatures(nn.Module):
+    EUPE_DIR = "/depot/mlp/data/robotvision/eupe" # Works on Gilbreth, Gautschi, Negishi, change if on Anvil
+    CHECKPOINT_PATH = f"{EUPE_DIR}/EUPE-ViT-S.pt"
+    SIZE = (768, 768) # pe_spatial_... is 512x512 I believe
+
+    def make_transform(self, resize_size: int = 256):
+        to_tensor = v2.ToImage()
+        resize = v2.Resize((resize_size, resize_size), antialias=True)
+        to_float = v2.ToDtype(torch.float32, scale=True)
+        normalize = v2.Normalize(
+            mean=(0.485, 0.456, 0.406),
+        std=(0.229, 0.224, 0.225),
+        )
+        return v2.Compose([to_tensor, resize, to_float, normalize])
+
+    def __init__(self, model_name: str = "EUPE-ViT-S", frozen: bool = True, feature_stage: int = -1):
+        super(EUPEFeatures, self).__init__()
+        if feature_stage != -1:
+            raise NotImplementedError
+        self.transform = self.make_transform(resize_size=self.SIZE[0]) # pics are like 900 x 1000
+        self.model = torch.hub.load(self.EUPE_DIR, "eupe_vits16", source="local", weights=self.CHECKPOINT_PATH)
+        if frozen:
+            for param in self.model.parameters():
+                param.requires_grad = False
+            self.model.eval()
+        self.dims = [384]  # feature dim for the last layer
+        self.patch_size = 16
+        self.n_tokens = (self.SIZE[0] // self.patch_size) ** 2 # 1024
+        self.data_config = {
+            "input_size": (3, self.SIZE[0], self.SIZE[1])
+        }
+
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        # x: (B, 3, H, W)
+        x_t = self.transform(x)  # preprocess
+        with torch.inference_mode():
+            with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
+                features = self.model.forward_features(x_t) # list of feature maps
+        clstoken, patchtokens = features["x_norm_clstoken"], features["x_norm_patchtokens"]
+        B, N, C = patchtokens.shape
+        H = W = int(N**0.5)
+        patchtokens = patchtokens.transpose(1, 2).reshape(B, C, H, W) # (B, C, H', W')
+        return [patchtokens]