Adds soft positional embedding with learnable projection.
| 18 | |
| 19 | |
| 20 | class SoftPositionEmbed(nn.Module): |
| 21 | """Adds soft positional embedding with learnable projection.""" |
| 22 | |
| 23 | def __init__(self, hidden_size, resolution): |
| 24 | """Builds the soft position embedding layer. |
| 25 | |
| 26 | Args: |
| 27 | hidden_size: Size of input feature dimension. |
| 28 | resolution: Tuple of integers specifying width and height of grid. |
| 29 | """ |
| 30 | super(SoftPositionEmbed, self).__init__() |
| 31 | self.proj = nn.Linear(4, hidden_size) |
| 32 | self.grid = build_grid(resolution) |
| 33 | |
| 34 | def forward(self, inputs): |
| 35 | return inputs + self.proj(self.grid) |
| 36 | |
| 37 | |
| 38 | def spatial_broadcast(slots, resolution): |