(self, *, H: int, W: int)
| 129 | self._init_weights() |
| 130 | |
| 131 | def forward(self, *, H: int, W: int) -> tuple[Tensor, Tensor]: |
| 132 | device = self.periods.device |
| 133 | dtype = self.dtype |
| 134 | dd = {"device": device, "dtype": dtype} |
| 135 | |
| 136 | # Prepare coords in range [-1, +1] |
| 137 | if self.normalize_coords == "max": |
| 138 | max_HW = max(H, W) |
| 139 | coords_h = torch.arange(0.5, H, **dd) / max_HW |
| 140 | coords_w = torch.arange(0.5, W, **dd) / max_HW |
| 141 | elif self.normalize_coords == "min": |
| 142 | min_HW = min(H, W) |
| 143 | coords_h = torch.arange(0.5, H, **dd) / min_HW |
| 144 | coords_w = torch.arange(0.5, W, **dd) / min_HW |
| 145 | elif self.normalize_coords == "separate": |
| 146 | coords_h = torch.arange(0.5, H, **dd) / H |
| 147 | coords_w = torch.arange(0.5, W, **dd) / W |
| 148 | else: |
| 149 | raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}") |
| 150 | coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1) |
| 151 | coords = coords.flatten(0, 1) |
| 152 | coords = 2.0 * coords - 1.0 # Shift range [0, 1] to [-1, +1] |
| 153 | |
| 154 | # Shift coords by adding a uniform value in [-shift, shift] |
| 155 | if self.training and self.shift_coords is not None: |
| 156 | shift_hw = torch.empty(2, **dd).uniform_(-self.shift_coords, self.shift_coords) |
| 157 | coords += shift_hw[None, :] |
| 158 | |
| 159 | # Jitter coords by multiplying the range [-1, 1] by a log-uniform value in [1/jitter, jitter] |
| 160 | if self.training and self.jitter_coords is not None: |
| 161 | jitter_max = np.log(self.jitter_coords) |
| 162 | jitter_min = -jitter_max |
| 163 | jitter_hw = torch.empty(2, **dd).uniform_(jitter_min, jitter_max).exp() |
| 164 | coords *= jitter_hw[None, :] |
| 165 | |
| 166 | # Rescale coords by multiplying the range [-1, 1] by a log-uniform value in [1/rescale, rescale] |
| 167 | if self.training and self.rescale_coords is not None: |
| 168 | rescale_max = np.log(self.rescale_coords) |
| 169 | rescale_min = -rescale_max |
| 170 | rescale_hw = torch.empty(1, **dd).uniform_(rescale_min, rescale_max).exp() |
| 171 | coords *= rescale_hw |
| 172 | |
| 173 | # Prepare angles and sin/cos |
| 174 | angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :] |
| 175 | angles = angles.flatten(1, 2) |
| 176 | angles = angles.tile(2) |
| 177 | cos = torch.cos(angles) |
| 178 | sin = torch.sin(angles) |
| 179 | |
| 180 | return (sin, cos) |
| 181 | |
| 182 | def _init_weights(self): |
| 183 | device = self.periods.device |
nothing calls this directly
no outgoing calls
no test coverage detected