| 247 | |
| 248 | |
| 249 | class DynRelPos2d(nn.Module): |
| 250 | |
| 251 | def __init__(self, embed_dim, num_heads, initial_value, heads_range): |
| 252 | """ |
| 253 | recurrent_chunk_size: (clh clw) |
| 254 | num_chunks: (nch ncw) |
| 255 | clh * clw == cl |
| 256 | nch * ncw == nc |
| 257 | |
| 258 | default: clh==clw, clh != clw is not implemented |
| 259 | """ |
| 260 | super().__init__() |
| 261 | angle = 1.0 / (10000 ** torch.linspace(0, 1, embed_dim // num_heads // 2)) |
| 262 | angle = angle.unsqueeze(-1).repeat(1, 2).flatten() |
| 263 | self.initial_value = initial_value |
| 264 | self.heads_range = heads_range |
| 265 | self.num_heads = num_heads |
| 266 | self.register_buffer('angle', angle) |
| 267 | |
| 268 | def generate_1d_decay(self, l: int, range_factor: Tensor): |
| 269 | """ |
| 270 | generate 1d decay mask, the result is l*l |
| 271 | """ |
| 272 | |
| 273 | range_factor=abs(range_factor) |
| 274 | |
| 275 | bs = range_factor.size(0) |
| 276 | ###print(f"Generating 1D decay mask for batch size: {bs} and length: {l}") |
| 277 | |
| 278 | heads_ranges = self.heads_range * torch.arange(self.num_heads, dtype=torch.float) / self.num_heads |
| 279 | heads_ranges = heads_ranges.to(range_factor.device) |
| 280 | ###print("Heads Ranges:") |
| 281 | ###print(heads_ranges) |
| 282 | ###print(f"Heads Ranges shape: {heads_ranges.shape}") |
| 283 | |
| 284 | range_factor = torch.sqrt(torch.sqrt(range_factor)) # (n) # give extra weight to smaller values |
| 285 | range_factor = range_factor[:, None] # (n 1) |
| 286 | ###print("Range Factor:") |
| 287 | ###print(range_factor) |
| 288 | ###print(f"Range Factor shape: {range_factor.shape}") |
| 289 | |
| 290 | ranges = (-self.initial_value - heads_ranges.repeat(bs, 1) * range_factor) |
| 291 | ###print("Ranges:") |
| 292 | decay = torch.log(1 - 2 ** ranges) # (b n) |
| 293 | ###print("Decay:") |
| 294 | ###print(decay) |
| 295 | ###print(f"Decay shape: {decay.shape}") |
| 296 | |
| 297 | index = torch.arange(l).to(decay) |
| 298 | ###print("Index:") |
| 299 | ###print(index) |
| 300 | ###print(f"Index shape: {index.shape}") |
| 301 | mask = index[:, None] - index[None, :] # (l l) |
| 302 | mask = mask.abs() # (l l) |
| 303 | # extend mask to batch size with one channel for each |
| 304 | mask = mask[None, None, :, :] # (1 1 l l) |
| 305 | ###print("Mask before decay application:") |
| 306 | ###print(mask) |