Merge pull request #9 from jxiong21029/simple-routing

Simplify router implementation

Author: flixmk

dit.py

@@ -243,8 +243,8 @@ def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor, **kwargs) -
         for idx, block in enumerate(self.blocks):
             if use_routing and idx == self.routes[route_count]['start_layer_idx']:
                 x_D_last = x.clone()
-                mask_info = self.router.get_mask(x, mask_ratio=self.routes[route_count]['selection_ratio'] if overwrite_selection_ratio is None else overwrite_selection_ratio)
-                x = self.router.start_route(x, mask_info)
+                ids_keep = self.router.get_mask(x, selection_rate=self.routes[route_count]['selection_ratio'] if overwrite_selection_ratio is None else overwrite_selection_ratio)
+                x = self.router.start_route(x, ids_keep)
 
             if fp32_next:
                 with torch.amp.autocast(device_type="cuda", enabled=False):
@@ -254,7 +254,7 @@ def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor, **kwargs) -
                 x = block(x, c)
 
             if use_routing and idx == self.routes[route_count]['end_layer_idx']:
-                x = self.router.end_route(x, mask_info, original_x=x_D_last)
+                x = self.router.end_route(x, ids_keep, original_x=x_D_last)
                 fp32_next = True
                 if route_count < len(self.routes) - 1:
                     route_count += 1

routing_module.py

@@ -11,35 +11,16 @@ def get_mask(self, x, selection_rate=0.0):
         num_keep = num_patches - num_mask
         noise_random = torch.rand(batch_size, num_patches, device=device)
         ids_shuffle = torch.argsort(noise_random, dim=1)
-        ids_restore = torch.argsort(ids_shuffle, dim=1)
         ids_keep = ids_shuffle[:, :num_keep]
-        ids_mask = ids_shuffle[:, num_keep:]
-        mask = torch.ones((batch_size, num_patches), device=device, dtype=torch.bool)
-        mask.scatter_(1, ids_keep, False)
-        return {
-            'mask': mask,           
-            'ids_keep': ids_keep,
-            'ids_mask': ids_mask,
-            'ids_shuffle': ids_shuffle,
-            'ids_restore': ids_restore
-        }
+        return ids_keep
 
-    def start_route(self, x, mask_info):
-        ids_shuffle = mask_info['ids_shuffle']
-        num_keep = mask_info['ids_keep'].size(1)
-        x_shuffled = x.gather(1, ids_shuffle.unsqueeze(-1).expand(-1, -1, x.size(2)))
-        masked_x = x_shuffled[:, :num_keep, :]
-        return masked_x
+    def start_route(self, x, ids_keep):
+        x_masked = x.gather(1, ids_keep.unsqueeze(-1).expand(-1, -1, x.size(2)))
+        return x_masked
 
-    def end_route(self, masked_x, mask_info, original_x):
-        batch_size, num_patches = mask_info['mask'].shape
-        num_keep = masked_x.size(1)
-        dim = masked_x.size(2)
-        device = masked_x.device
-        ids_restore = mask_info['ids_restore']
-        x_unshuffled = torch.empty((batch_size, num_patches, dim), device=device)
-        x_unshuffled[:, :num_keep, :] = masked_x
-        x_shuffled = original_x.gather(1, mask_info['ids_shuffle'].unsqueeze(-1).expand(-1, -1, dim))
-        x_unshuffled[:, num_keep:, :] = x_shuffled[:, num_keep:, :]
-        x_unmasked = x_unshuffled.gather(1, ids_restore.unsqueeze(-1).expand(-1, -1, dim))
-        return x_unmasked
+    def end_route(self, masked_x, ids_keep, original_x):
+        # (jerry) scatter is out-of-place, so this is safe
+        x_unmasked = original_x.scatter(
+            1, ids_keep.unsqueeze(-1).expand(-1, -1, original_x.size(2)), masked_x
+        )
+        return x_unmasked