Store a single block error into the matching bucket. Args: error_block: (block_size, C, H, W) tensor timestep_index: int, raw index in [0, num_train_timesteps) block_pos: int, global block position; required iff num_blocks>0
(self, error_block, timestep_index, block_pos=None)
| 109 | |
| 110 | # ------------------------------------------------------------------ add |
| 111 | def add(self, error_block, timestep_index, block_pos=None): |
| 112 | """Store a single block error into the matching bucket. |
| 113 | |
| 114 | Args: |
| 115 | error_block: (block_size, C, H, W) tensor |
| 116 | timestep_index: int, raw index in [0, num_train_timesteps) |
| 117 | block_pos: int, global block position; required iff num_blocks>0 |
| 118 | """ |
| 119 | t = self._t_bucket(timestep_index) |
| 120 | if not self._is_owned_t(t): |
| 121 | return |
| 122 | key = self._make_key(t, block_pos) |
| 123 | # Store in the source dtype on CPU to match SVI (which keeps bf16), |
| 124 | # cutting buffer memory in half vs. casting to fp32. |
| 125 | entry = error_block.detach().to("cpu", copy=True) |
| 126 | |
| 127 | buf = self.buckets[key] |
| 128 | if len(buf) < self.max_size: |
| 129 | buf.append(entry) |
| 130 | else: |
| 131 | if self.replacement_strategy == "fifo": |
| 132 | buf.pop(0) |
| 133 | buf.append(entry) |
| 134 | elif self.replacement_strategy == "l2": |
| 135 | stacked = torch.stack(buf) |
| 136 | dists = (stacked - entry.unsqueeze(0)).flatten(1).norm(dim=1) |
| 137 | most_similar = torch.argmin(dists).item() |
| 138 | buf[most_similar] = entry |
| 139 | else: # "random" (default) |
| 140 | idx = random.randint(0, self.max_size - 1) |
| 141 | buf[idx] = entry |
| 142 | self.total_added += 1 |
| 143 | |
| 144 | # ------------------------------------------------------------------ sample |
| 145 | def sample(self, timestep_index, device, dtype, block_pos=None): |
no test coverage detected