| import torch | |
| import math | |
| from .weights import RegionModel | |
| from .layers import linear, mlp | |
| def fourier_features(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Applies Fourier feature mapping to input tensor x using frequency matrix w. This | |
| projects inputs through sinusoidal functions to create higher dimensional features | |
| that help mitigate spectral bias - the tendency of neural networks to learn | |
| low-frequency functions more easily than high-frequency ones. By explicitly | |
| mapping inputs to higher frequencies through sin/cos transformations, we enable | |
| better learning of fine details and higher frequency patterns. | |
| Args: | |
| x: Input tensor to transform | |
| w: Matrix of frequencies for the Fourier features transformation | |
| Returns: | |
| Concatenated cosine and sine transformed features as a tensor | |
| """ | |
| f = 2 * math.pi * x @ w | |
| return torch.cat([f.cos(), f.sin()], dim=-1) | |
| def encode_coordinate(coord: torch.Tensor, w: RegionModel) -> torch.Tensor: | |
| """ | |
| Takes as input a tensor containing a single float coordinate value (x or y) | |
| and encodes it into hidden states for input to the text model. | |
| Args: | |
| coord: Tensor with single float coordinate value | |
| Returns: | |
| Encoded hidden states tensor for input to text model | |
| """ | |
| return linear(fourier_features(coord, w.coord_features), w.coord_encoder) | |
| def decode_coordinate(hidden_state: torch.Tensor, w: RegionModel) -> torch.Tensor: | |
| """ | |
| Takes as input the last hidden state from the text model and outputs a single logit | |
| representing either an x or y coordinate prediction. | |
| Args: | |
| hidden_state: The final hidden state tensor from the text model. | |
| Returns: | |
| A single logit representing the predicted coordinate value (x or y) | |
| """ | |
| return mlp(hidden_state, w.coord_decoder) | |
| def encode_size(size: torch.Tensor, w: RegionModel) -> torch.Tensor: | |
| """ | |
| Takes a tensor containing normalized width and height values in range [0,1] | |
| and encodes them into hidden states for input to the text model. | |
| Args: | |
| size: Tensor with two floats for width and height in range [0,1] | |
| Returns: | |
| Encoded hidden states tensor for input to text model | |
| """ | |
| return linear(fourier_features(size, w.size_features), w.size_encoder) | |
| def decode_size(hidden_state: torch.Tensor, w: RegionModel) -> torch.Tensor: | |
| """ | |
| Takes as input the last hidden state from the text model and outputs two logits | |
| for width and height respectively. | |
| Args: | |
| hidden_state: The final hidden state tensor from the text model. | |
| Returns: | |
| A tensor containing two logits - one for predicted width and one for | |
| predicted height. | |
| """ | |
| return mlp(hidden_state, w.size_decoder).view(2, -1) | |