This is a sinusoidal position encoding that generalized to 2-dimensional images
| 36 | return x + self.pe[:, :, :x.size(2), :x.size(3)] |
| 37 | |
| 38 | class LinearPositionEncoding(nn.Module): |
| 39 | """ |
| 40 | This is a sinusoidal position encoding that generalized to 2-dimensional images |
| 41 | """ |
| 42 | |
| 43 | def __init__(self, d_model, max_shape=(256, 256)): |
| 44 | """ |
| 45 | Args: |
| 46 | max_shape (tuple): for 1/8 featmap, the max length of 256 corresponds to 2048 pixels |
| 47 | """ |
| 48 | super().__init__() |
| 49 | |
| 50 | pe = torch.zeros((d_model, *max_shape)) |
| 51 | y_position = (torch.ones(max_shape).cumsum(0).float().unsqueeze(0) - 1) / max_shape[0] |
| 52 | x_position = (torch.ones(max_shape).cumsum(1).float().unsqueeze(0) - 1) / max_shape[1] |
| 53 | div_term = torch.arange(0, d_model//2, 2).float() |
| 54 | div_term = div_term[:, None, None] # [C//4, 1, 1] |
| 55 | pe[0::4, :, :] = torch.sin(x_position * div_term * math.pi) |
| 56 | pe[1::4, :, :] = torch.cos(x_position * div_term * math.pi) |
| 57 | pe[2::4, :, :] = torch.sin(y_position * div_term * math.pi) |
| 58 | pe[3::4, :, :] = torch.cos(y_position * div_term * math.pi) |
| 59 | |
| 60 | self.register_buffer('pe', pe.unsqueeze(0), persistent=False) # [1, C, H, W] |
| 61 | |
| 62 | def forward(self, x): |
| 63 | """ |
| 64 | Args: |
| 65 | x: [N, C, H, W] |
| 66 | """ |
| 67 | # assert x.shape[2] == 80 and x.shape[3] == 80 |
| 68 | |
| 69 | return x + self.pe[:, :, :x.size(2), :x.size(3)] |
| 70 | |
| 71 | class LearnedPositionEncoding(nn.Module): |
| 72 | """ |
nothing calls this directly
no outgoing calls
no test coverage detected