Args: max_shape (tuple): for 1/8 featmap, the max length of 256 corresponds to 2048 pixels
(self, d_model, max_shape=(256, 256))
| 10 | """ |
| 11 | |
| 12 | def __init__(self, d_model, max_shape=(256, 256)): |
| 13 | """ |
| 14 | Args: |
| 15 | max_shape (tuple): for 1/8 featmap, the max length of 256 corresponds to 2048 pixels |
| 16 | """ |
| 17 | super().__init__() |
| 18 | |
| 19 | pe = torch.zeros((d_model, *max_shape)) |
| 20 | y_position = torch.ones(max_shape).cumsum(0).float().unsqueeze(0) |
| 21 | x_position = torch.ones(max_shape).cumsum(1).float().unsqueeze(0) |
| 22 | div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / d_model//2)) |
| 23 | div_term = div_term[:, None, None] # [C//4, 1, 1] |
| 24 | pe[0::4, :, :] = torch.sin(x_position * div_term) |
| 25 | pe[1::4, :, :] = torch.cos(x_position * div_term) |
| 26 | pe[2::4, :, :] = torch.sin(y_position * div_term) |
| 27 | pe[3::4, :, :] = torch.cos(y_position * div_term) |
| 28 | |
| 29 | self.register_buffer('pe', pe.unsqueeze(0)) # [1, C, H, W] |
| 30 | |
| 31 | def forward(self, x): |
| 32 | """ |