| """Masking strategy for Grid-JEPA.""" |
| import random |
| import torch |
|
|
| class GridMaskingStrategy: |
| def __init__(self, grid_size=64, num_target_blocks=4, |
| target_scale=(0.05, 0.15), target_aspect_ratio=(0.5, 2.0), |
| context_scale=(0.85, 1.0), min_target_size=2): |
| self.grid_size = grid_size |
| self.num_patches = grid_size * grid_size |
| self.num_target_blocks = num_target_blocks |
| self.target_scale = target_scale |
| self.target_aspect_ratio = target_aspect_ratio |
| self.context_scale = context_scale |
| self.min_target_size = min_target_size |
| |
| def _sample_block(self, scale_range, aspect_range): |
| total = self.grid_size * self.grid_size |
| area = max(int(random.uniform(*scale_range) * total), self.min_target_size * self.min_target_size) |
| aspect = random.uniform(*aspect_range) |
| h = max(min(int((area / aspect) ** 0.5), self.grid_size), self.min_target_size) |
| w = max(min(int(area / h), self.grid_size), self.min_target_size) |
| rs = random.randint(0, max(0, self.grid_size - h)) |
| cs = random.randint(0, max(0, self.grid_size - w)) |
| return rs, min(rs + h, self.grid_size), cs, min(cs + w, self.grid_size) |
| |
| def _block_to_indices(self, block): |
| rs, re, cs, ce = block |
| return [r * self.grid_size + c for r in range(rs, re) for c in range(cs, ce)] |
| |
| def sample_masks(self): |
| target_blocks = [self._sample_block(self.target_scale, self.target_aspect_ratio) for _ in range(self.num_target_blocks)] |
| target_indices = set() |
| for b in target_blocks: |
| target_indices.update(self._block_to_indices(b)) |
| target_indices = sorted(list(target_indices)) |
| ctx_block = self._sample_block(self.context_scale, (0.75, 1.5)) |
| ctx_indices = [i for i in self._block_to_indices(ctx_block) if i not in target_indices] |
| tgt_mask = torch.zeros(self.num_patches, dtype=torch.bool) |
| tgt_mask[target_indices] = True |
| ctx_mask = torch.zeros(self.num_patches, dtype=torch.bool) |
| ctx_mask[ctx_indices] = True |
| return { |
| 'target_indices': torch.tensor(target_indices, dtype=torch.long), |
| 'context_indices': torch.tensor(ctx_indices, dtype=torch.long), |
| 'target_mask': tgt_mask, |
| 'context_mask': ctx_mask, |
| } |
| |
| def sample_masks_batch(self, batch_size): |
| all_ti, all_ci, all_tm, all_cm = [], [], [], [] |
| for _ in range(batch_size): |
| m = self.sample_masks() |
| all_ti.append(m['target_indices']) |
| all_ci.append(m['context_indices']) |
| all_tm.append(m['target_mask']) |
| all_cm.append(m['context_mask']) |
| max_t = max(len(t) for t in all_ti) |
| max_c = max(len(c) for c in all_ci) |
| ti_p = torch.zeros(batch_size, max_t, dtype=torch.long) |
| ci_p = torch.zeros(batch_size, max_c, dtype=torch.long) |
| for i, (t, c) in enumerate(zip(all_ti, all_ci)): |
| ti_p[i, :len(t)] = t |
| ci_p[i, :len(c)] = c |
| return { |
| 'target_indices': ti_p, |
| 'target_lengths': torch.tensor([len(t) for t in all_ti]), |
| 'context_indices': ci_p, |
| 'context_lengths': torch.tensor([len(c) for c in all_ci]), |
| 'target_mask': torch.stack(all_tm), |
| 'context_mask': torch.stack(all_cm), |
| } |
|
|
