Skip to content

rslearn.models.tokens_to_channels

tokens_to_channels

Project each token of a TokenFeatureMaps and stack the results as channels.

TokensToChannels

Bases: IntermediateComponent

Apply a shared linear projection per token and lay the tokens out as channels.

Given a BCHWN token feature map, a linear layer C -> out_dim is applied to every token independently and the results are concatenated along the channel dimension in token order, producing a B x (N * out_dim) x H x W feature map.

With out_dim=1 this turns per-timestep tokens into one logit per timestep at each location, e.g. for predicting the timestep at which an event occurred as an N-way classification per pixel.

If the input TokenFeatureMaps has masks, the output channels corresponding to invalid tokens are set to mask_fill_value (the output shape does not change).

Source code in rslearn/models/tokens_to_channels.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class TokensToChannels(IntermediateComponent):
    """Apply a shared linear projection per token and lay the tokens out as channels.

    Given a BCHWN token feature map, a linear layer C -> out_dim is applied to every
    token independently and the results are concatenated along the channel dimension
    in token order, producing a B x (N * out_dim) x H x W feature map.

    With out_dim=1 this turns per-timestep tokens into one logit per timestep at each
    location, e.g. for predicting the timestep at which an event occurred as an
    N-way classification per pixel.

    If the input TokenFeatureMaps has masks, the output channels corresponding to
    invalid tokens are set to mask_fill_value (the output shape does not change).
    """

    def __init__(
        self, in_dim: int, out_dim: int = 1, mask_fill_value: float = 0.0
    ) -> None:
        """Create a new TokensToChannels.

        Args:
            in_dim: the token embedding dimension C.
            out_dim: the number of output values per token.
            mask_fill_value: the value to write into the output channels of tokens
                that are marked invalid by the input mask (if any).
        """
        super().__init__()
        self.out_dim = out_dim
        self.mask_fill_value = mask_fill_value
        self.linear = nn.Linear(in_dim, out_dim)

    def forward_for_map(
        self, feat_tokens: torch.Tensor, mask: torch.Tensor | None = None
    ) -> torch.Tensor:
        """Project a single BCHWN token feature map to B(N*out_dim)HW.

        Args:
            feat_tokens: the BCHWN token feature map.
            mask: optional BHWN bool mask of valid tokens. Outputs for invalid tokens
                are set to mask_fill_value.
        """
        x = rearrange(feat_tokens, "b c h w n -> b h w n c")
        x = self.linear(x)  # (B, H, W, N, out_dim)
        if mask is not None:
            x = x.masked_fill(~mask.unsqueeze(-1), self.mask_fill_value)
        return rearrange(x, "b h w n d -> b (n d) h w")

    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
        """Project the tokens in each feature map.

        Args:
            intermediates: the output from the previous component, which must be a
                TokenFeatureMaps.
            context: the model context.

        Returns:
            a FeatureMaps where each map is B x (N * out_dim) x H x W.
        """
        if not isinstance(intermediates, TokenFeatureMaps):
            raise ValueError("input to TokensToChannels must be a TokenFeatureMaps")
        return FeatureMaps(
            [
                self.forward_for_map(feat, mask)
                for feat, mask in zip(
                    intermediates.feature_maps, intermediates.get_masks()
                )
            ]
        )

forward_for_map 

forward_for_map(feat_tokens: Tensor, mask: Tensor | None = None) -> Tensor

Project a single BCHWN token feature map to B(N*out_dim)HW.

Parameters:

Name Type Description Default
feat_tokens Tensor

the BCHWN token feature map.

 required
mask Tensor | None

optional BHWN bool mask of valid tokens. Outputs for invalid tokens are set to mask_fill_value.

 None
Source code in rslearn/models/tokens_to_channels.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def forward_for_map(
    self, feat_tokens: torch.Tensor, mask: torch.Tensor | None = None
) -> torch.Tensor:
    """Project a single BCHWN token feature map to B(N*out_dim)HW.

    Args:
        feat_tokens: the BCHWN token feature map.
        mask: optional BHWN bool mask of valid tokens. Outputs for invalid tokens
            are set to mask_fill_value.
    """
    x = rearrange(feat_tokens, "b c h w n -> b h w n c")
    x = self.linear(x)  # (B, H, W, N, out_dim)
    if mask is not None:
        x = x.masked_fill(~mask.unsqueeze(-1), self.mask_fill_value)
    return rearrange(x, "b h w n d -> b (n d) h w")

forward 

forward(intermediates: Any, context: ModelContext) -> FeatureMaps

Project the tokens in each feature map.

Parameters:

Name Type Description Default
intermediates Any

the output from the previous component, which must be a TokenFeatureMaps.

 required
context ModelContext

the model context.

 required

Returns:

Type Description
FeatureMaps

a FeatureMaps where each map is B x (N * out_dim) x H x W.
Source code in rslearn/models/tokens_to_channels.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
    """Project the tokens in each feature map.

    Args:
        intermediates: the output from the previous component, which must be a
            TokenFeatureMaps.
        context: the model context.

    Returns:
        a FeatureMaps where each map is B x (N * out_dim) x H x W.
    """
    if not isinstance(intermediates, TokenFeatureMaps):
        raise ValueError("input to TokensToChannels must be a TokenFeatureMaps")
    return FeatureMaps(
        [
            self.forward_for_map(feat, mask)
            for feat, mask in zip(
                intermediates.feature_maps, intermediates.get_masks()
            )
        ]
    )