This is a sinusoidal position encoding that generalized to 2-dimensional images
| 69 | return x + self.pe[:, :, :x.size(2), :x.size(3)] |
| 70 | |
| 71 | class LearnedPositionEncoding(nn.Module): |
| 72 | """ |
| 73 | This is a sinusoidal position encoding that generalized to 2-dimensional images |
| 74 | """ |
| 75 | |
| 76 | def __init__(self, d_model, max_shape=(80, 80)): |
| 77 | """ |
| 78 | Args: |
| 79 | max_shape (tuple): for 1/8 featmap, the max length of 256 corresponds to 2048 pixels |
| 80 | """ |
| 81 | super().__init__() |
| 82 | |
| 83 | self.pe = nn.Parameter(torch.randn(1, max_shape[0], max_shape[1], d_model)) |
| 84 | |
| 85 | def forward(self, x): |
| 86 | """ |
| 87 | Args: |
| 88 | x: [N, C, H, W] |
| 89 | """ |
| 90 | # assert x.shape[2] == 80 and x.shape[3] == 80 |
| 91 | |
| 92 | return x + self.pe |
nothing calls this directly
no outgoing calls
no test coverage detected