| 457 | |
| 458 | |
| 459 | class PositionalEncoding3D(nn.Module): |
| 460 | def __init__(self, channels): |
| 461 | """ |
| 462 | :param channels: The last dimension of the tensor you want to apply pos emb to. |
| 463 | """ |
| 464 | self.orig_ch = channels |
| 465 | super(PositionalEncoding3D, self).__init__() |
| 466 | channels = int(np.ceil(channels / 6) * 2) |
| 467 | if channels % 2: |
| 468 | channels += 1 |
| 469 | self.channels = channels |
| 470 | inv_freq = 1.0 / (10000 ** (torch.arange(0, channels, 2).float() / channels)) |
| 471 | self.register_buffer("inv_freq", inv_freq) |
| 472 | |
| 473 | def forward(self, tensor, input_range=None): |
| 474 | """ |
| 475 | :param tensor: A 5d tensor of size (batch_size, x, y, z, ch) |
| 476 | :return: Positional Encoding Matrix of size (batch_size, x, y, z, ch) |
| 477 | """ |
| 478 | pos_x, pos_y, pos_z = tensor[:, :, 0], tensor[:, :, 1], tensor[:, :, 2] |
| 479 | sin_inp_x = torch.einsum("bi,j->bij", pos_x, self.inv_freq) |
| 480 | sin_inp_y = torch.einsum("bi,j->bij", pos_y, self.inv_freq) |
| 481 | sin_inp_z = torch.einsum("bi,j->bij", pos_z, self.inv_freq) |
| 482 | emb_x = torch.cat((sin_inp_x.sin(), sin_inp_x.cos()), dim=-1) |
| 483 | |
| 484 | emb_y = torch.cat((sin_inp_y.sin(), sin_inp_y.cos()), dim=-1) |
| 485 | emb_z = torch.cat((sin_inp_z.sin(), sin_inp_z.cos()), dim=-1) |
| 486 | |
| 487 | emb = torch.cat((emb_x, emb_y, emb_z), dim=-1) |
| 488 | return emb[:, :, :self.orig_ch].permute((0, 2, 1)) |
| 489 | |
| 490 | |
| 491 | class SelfAttentionLayer(nn.Module): |