Slice DeMemWM diffusion updates to target frames

algorithms/dememwm/models/diffusion.py

@@ -188,7 +188,13 @@ class Diffusion(nn.Module):
         model_output = self.model(x, t, action_cond, **model_kwargs)
         model_output = model_output.permute(1,0,2,3,4)
         x = x.permute(1,0,2,3,4)
-        t = t.permute(1,0)        
+        t = t.permute(1,0)
+
+        target_frames = self._target_frames(x, frame_memory_segments)
+        if frame_memory_segments is not None:
+            model_output = model_output[:target_frames]
+            x = x[:target_frames]
+            t = t[:target_frames]
 
         if self.objective == "pred_noise":
             pred_noise = torch.clamp(model_output, -self.clip_noise, self.clip_noise)
@@ -262,6 +268,10 @@ class Diffusion(nn.Module):
             frame_memory_masks=frame_memory_masks,
             frame_memory_pose=frame_memory_pose, image_hw=image_hw)
         x_start = model_pred.pred_x_start
+        target_frames = self._target_frames(x, frame_memory_segments)
+        if frame_memory_segments is not None:
+            x = x[:target_frames]
+            t = t[:target_frames]
         return self.q_posterior(x_start=x_start, x_t=x, t=t)
 
     def compute_loss_weights(self, noise_levels: torch.Tensor):
@@ -328,19 +338,27 @@ class Diffusion(nn.Module):
 
         pred = model_pred.model_out
         x_pred = model_pred.pred_x_start
+        target_frames = self._target_frames(x, frame_memory_segments)
+        target_x = x
+        target_noise = noise
+        target_noise_levels = noise_levels
+        if frame_memory_segments is not None:
+            target_x = x[:target_frames]
+            target_noise = noise[:target_frames]
+            target_noise_levels = noise_levels[:target_frames]
 
         if self.objective == "pred_noise":
-            target = noise
+            target = target_noise
         elif self.objective == "pred_x0":
-            target = x
+            target = target_x
         elif self.objective == "pred_v":
-            target = self.predict_v(x, noise_levels, noise)
+            target = self.predict_v(target_x, target_noise_levels, target_noise)
         else:
             raise ValueError(f"unknown objective {self.objective}")
 
         # 训练的时候每个frame随便给噪声
         loss = F.mse_loss(pred, target.detach(), reduction="none")
-        loss_weight = self.compute_loss_weights(noise_levels)
+        loss_weight = self.compute_loss_weights(target_noise_levels)
 
         loss_weight = loss_weight.view(*loss_weight.shape, *((1,) * (loss.ndim - 2)))
 
@@ -458,16 +476,25 @@ class Diffusion(nn.Module):
                 image_hw=image_hw
             )
 
+        target_frames = self._target_frames(x, frame_memory_segments)
+        update_curr_noise_level = curr_noise_level
+        update_clipped_curr_noise_level = clipped_curr_noise_level
+        update_orig_x = orig_x
+        if frame_memory_segments is not None:
+            update_curr_noise_level = curr_noise_level[:target_frames]
+            update_clipped_curr_noise_level = clipped_curr_noise_level[:target_frames]
+            update_orig_x = orig_x[:target_frames]
+
         noise = torch.where(
-            self.add_shape_channels(clipped_curr_noise_level > 0),
-            torch.randn_like(x),
+            self.add_shape_channels(update_clipped_curr_noise_level > 0),
+            torch.randn_like(model_mean),
             0,
         )
         noise = torch.clamp(noise, -self.clip_noise, self.clip_noise)
         x_pred = model_mean + torch.exp(0.5 * model_log_variance) * noise
 
         # only update frames where the noise level decreases
-        return torch.where(self.add_shape_channels(curr_noise_level == -1), orig_x, x_pred)
+        return torch.where(self.add_shape_channels(update_curr_noise_level == -1), update_orig_x, x_pred)
 
     def ddim_sample_step(
         self,
@@ -502,15 +529,26 @@ class Diffusion(nn.Module):
         )
         x = torch.where(self.add_shape_channels(curr_noise_level < 0), scaled_context, orig_x)
 
-        alpha = self.alphas_cumprod[clipped_curr_noise_level]
+        target_frames = self._target_frames(x, frame_memory_segments)
+        update_curr_noise_level = curr_noise_level
+        update_clipped_curr_noise_level = clipped_curr_noise_level
+        update_next_noise_level = next_noise_level
+        update_orig_x = orig_x
+        if frame_memory_segments is not None:
+            update_curr_noise_level = curr_noise_level[:target_frames]
+            update_clipped_curr_noise_level = clipped_curr_noise_level[:target_frames]
+            update_next_noise_level = next_noise_level[:target_frames]
+            update_orig_x = orig_x[:target_frames]
+
+        alpha = self.alphas_cumprod[update_clipped_curr_noise_level]
         alpha_next = torch.where(
-            next_noise_level < 0,
-            torch.ones_like(next_noise_level),
-            self.alphas_cumprod[next_noise_level],
+            update_next_noise_level < 0,
+            torch.ones_like(update_next_noise_level),
+            self.alphas_cumprod[update_next_noise_level],
         )
         sigma = torch.where(
-            next_noise_level < 0,
-            torch.zeros_like(next_noise_level),
+            update_next_noise_level < 0,
+            torch.zeros_like(update_next_noise_level),
             self.ddim_sampling_eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt(),
         )
         c = (1 - alpha_next - sigma**2).sqrt()
@@ -522,6 +560,9 @@ class Diffusion(nn.Module):
         if guidance_fn is not None:
             with torch.enable_grad():
                 x = x.detach().requires_grad_()
+                update_x = x
+                if frame_memory_segments is not None:
+                    update_x = x[:target_frames]
 
                 model_pred = self.model_predictions(
                     x=x,
@@ -543,9 +584,11 @@ class Diffusion(nn.Module):
                     guidance_loss,
                     x,
                 )[0]
+                if frame_memory_segments is not None:
+                    grad = grad[:target_frames]
 
                 pred_noise = model_pred.pred_noise + (1 - alpha_next).sqrt() * grad
-                x_start = self.predict_start_from_noise(x, clipped_curr_noise_level, pred_noise)
+                x_start = self.predict_start_from_noise(update_x, update_clipped_curr_noise_level, pred_noise)
 
         else:
             # print(clipped_curr_noise_level)
@@ -565,17 +608,20 @@ class Diffusion(nn.Module):
             )
             x_start = model_pred.pred_x_start
             pred_noise = model_pred.pred_noise
+            update_x = x
+            if frame_memory_segments is not None:
+                update_x = x[:target_frames]
 
-        noise = torch.randn_like(x)
+        noise = torch.randn_like(update_x)
         noise = torch.clamp(noise, -self.clip_noise, self.clip_noise)
 
         x_pred = x_start * alpha_next.sqrt() + pred_noise * c + sigma * noise
 
         # only update frames where the noise level decreases
-        mask = curr_noise_level == next_noise_level
+        mask = update_curr_noise_level == update_next_noise_level
         x_pred = torch.where(
             self.add_shape_channels(mask),
-            orig_x,
+            update_orig_x,
             x_pred,
         )
 

tests/test_dememwm_diffusion.py

@@ -0,0 +1,155 @@
+import unittest
+
+import torch
+from torch import nn
+
+from algorithms.dememwm.models.diffusion import Diffusion
+
+
+class FakeDenoiser(nn.Module):
+    def __init__(self, output_frames=None):
+        super().__init__()
+        self.output_frames = output_frames
+        self.calls = []
+
+    def forward(self, x, t, action_cond, **kwargs):
+        self.calls.append({
+            "x_shape": tuple(x.shape),
+            "t_shape": tuple(t.shape),
+            "kwargs": kwargs,
+        })
+        frames = self.output_frames if self.output_frames is not None else x.shape[1]
+        return torch.zeros((x.shape[0], frames, *x.shape[2:]), device=x.device, dtype=x.dtype)
+
+
+def _make_diffusion(output_frames=None):
+    diffusion = Diffusion.__new__(Diffusion)
+    nn.Module.__init__(diffusion)
+    diffusion.x_shape = torch.Size((1, 1, 1))
+    diffusion.timesteps = 4
+    diffusion.sampling_timesteps = 4
+    diffusion.is_ddim_sampling = False
+    diffusion.objective = "pred_noise"
+    diffusion.use_fused_snr = False
+    diffusion.snr_clip = 5.0
+    diffusion.cum_snr_decay = 0.9
+    diffusion.ddim_sampling_eta = 0.0
+    diffusion.clip_noise = 10.0
+    diffusion.stabilization_level = 1
+    diffusion.model = FakeDenoiser(output_frames=output_frames)
+
+    betas = torch.tensor([0.05, 0.10, 0.15, 0.20], dtype=torch.float32)
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_cumprod_prev = torch.nn.functional.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
+    posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+    snr = alphas_cumprod / (1.0 - alphas_cumprod)
+
+    diffusion.register_buffer("betas", betas)
+    diffusion.register_buffer("alphas_cumprod", alphas_cumprod)
+    diffusion.register_buffer("alphas_cumprod_prev", alphas_cumprod_prev)
+    diffusion.register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod))
+    diffusion.register_buffer("sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod))
+    diffusion.register_buffer("log_one_minus_alphas_cumprod", torch.log(1.0 - alphas_cumprod))
+    diffusion.register_buffer("sqrt_recip_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod))
+    diffusion.register_buffer("sqrt_recipm1_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod - 1.0))
+    diffusion.register_buffer("posterior_variance", posterior_variance)
+    diffusion.register_buffer("posterior_log_variance_clipped", torch.log(posterior_variance.clamp(min=1e-20)))
+    diffusion.register_buffer("posterior_mean_coef1", betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod))
+    diffusion.register_buffer("posterior_mean_coef2", (1.0 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - alphas_cumprod))
+    diffusion.register_buffer("snr", snr)
+    diffusion.register_buffer("clipped_snr", snr.clamp(max=diffusion.snr_clip))
+    return diffusion
+
+
+def _packed_inputs():
+    x = torch.arange(5, dtype=torch.float32).view(5, 1, 1, 1, 1)
+    action_cond = torch.zeros((5, 1, 3), dtype=torch.float32)
+    noise_levels = torch.tensor([[1], [2], [0], [0], [0]], dtype=torch.long)
+    segments = {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1}
+    return x, action_cond, noise_levels, segments
+
+
+class DeMemWMDiffusionTargetOnlyTests(unittest.TestCase):
+    def test_forward_noises_packed_input_but_returns_target_loss(self):
+        torch.manual_seed(0)
+        diffusion = _make_diffusion(output_frames=2)
+        x, action_cond, noise_levels, segments = _packed_inputs()
+
+        x_pred, loss = diffusion(
+            x,
+            action_cond,
+            None,
+            noise_levels=noise_levels,
+            reference_length=0,
+            frame_memory_segments=segments,
+        )
+
+        self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1))
+        self.assertEqual(tuple(x_pred.shape), (2, 1, 1, 1, 1))
+        self.assertEqual(tuple(loss.shape), (2, 1, 1, 1, 1))
+
+    def test_forward_without_frame_memory_keeps_full_length(self):
+        torch.manual_seed(0)
+        diffusion = _make_diffusion()
+        x, action_cond, noise_levels, _ = _packed_inputs()
+
+        x_pred, loss = diffusion(
+            x,
+            action_cond,
+            None,
+            noise_levels=noise_levels,
+            reference_length=0,
+        )
+
+        self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1))
+        self.assertEqual(tuple(x_pred.shape), tuple(x.shape))
+        self.assertEqual(tuple(loss.shape), tuple(x.shape))
+
+    def test_posterior_and_sample_steps_return_target_frames_for_frame_memory(self):
+        torch.manual_seed(0)
+        x, action_cond, _, segments = _packed_inputs()
+        curr = torch.tensor([[2], [1], [-1], [-1], [-1]], dtype=torch.long)
+        next_level = torch.tensor([[1], [0], [-1], [-1], [-1]], dtype=torch.long)
+
+        diffusion = _make_diffusion(output_frames=2)
+        mean, variance, log_variance = diffusion.p_mean_variance(
+            x,
+            curr,
+            action_cond=action_cond,
+            pose_cond=None,
+            reference_length=0,
+            frame_memory_segments=segments,
+        )
+        self.assertEqual(tuple(mean.shape), (2, 1, 1, 1, 1))
+        self.assertEqual(tuple(variance.shape), (2, 1, 1, 1, 1))
+        self.assertEqual(tuple(log_variance.shape), (2, 1, 1, 1, 1))
+
+        diffusion = _make_diffusion(output_frames=2)
+        ddpm = diffusion.ddpm_sample_step(
+            x,
+            action_cond,
+            None,
+            curr_noise_level=curr,
+            reference_length=0,
+            frame_memory_segments=segments,
+        )
+        self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1))
+        self.assertEqual(tuple(ddpm.shape), (2, 1, 1, 1, 1))
+
+        diffusion = _make_diffusion(output_frames=2)
+        ddim = diffusion.ddim_sample_step(
+            x,
+            action_cond,
+            None,
+            curr_noise_level=curr,
+            next_noise_level=next_level,
+            reference_length=0,
+            frame_memory_segments=segments,
+        )
+        self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1))
+        self.assertEqual(tuple(ddim.shape), (2, 1, 1, 1, 1))
+
+
+if __name__ == "__main__":
+    unittest.main()