Skip to content

rslearn.models.task_embedding

task_embedding

Task embedding modules.

PositionalEncoding

Bases: Module

Simple sinusoidal positional encoding for the task embedding. From torch docs.

Source code in rslearn/models/task_embedding.py 
10
11
12
13
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
class PositionalEncoding(nn.Module):
    """Simple sinusoidal positional encoding for the task embedding. From torch docs."""

    def __init__(self, d_model: int, dropout: float = 0.0, max_len: int = 1024):
        """Initialize the positional encoding module.

        Args:
            d_model: The dimension of the model.
            dropout: The dropout rate.
            max_len: The maximum length of the sequence.
        """
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
        )
        pe = torch.zeros(max_len, 1, d_model)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        pe[:, 0, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply positional encoding to the input tensor.

        Args:
            x: Tensor, shape ``[seq_len, batch_size, embedding_dim]``
        """
        x = x + self.pe[: x.size(0)]
        return self.dropout(x)

forward 

forward(x: Tensor) -> Tensor

Apply positional encoding to the input tensor.

Parameters:

Name Type Description Default
x Tensor

Tensor, shape [seq_len, batch_size, embedding_dim]

 required
Source code in rslearn/models/task_embedding.py 
33
34
35
36
37
38
39
40
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Apply positional encoding to the input tensor.

    Args:
        x: Tensor, shape ``[seq_len, batch_size, embedding_dim]``
    """
    x = x + self.pe[: x.size(0)]
    return self.dropout(x)

BaseTaskEmbedding

Bases: Module

Base class for task embedding modules.

Source code in rslearn/models/task_embedding.py 
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
class BaseTaskEmbedding(torch.nn.Module):
    """Base class for task embedding modules."""

    def __init__(self, encoder_embedding_size: int) -> None:
        """Initialize the base task embedding module.

        Args:
            encoder_embedding_size: The size of the encoder embedding.
        """
        super().__init__()
        self.encoder_embedding_size = encoder_embedding_size

    def register_tasks(self, task_names: list[str]) -> None:
        """Register the tasks.

        This must happen post-init so that we can dynamically determine
        the tasks to use, so it doesn't have to be specified in the config.

        Args:
            task_names: The names of the tasks.
        """
        raise NotImplementedError

    def compute_embeds(
        self,
        features: list[torch.tensor],
        inputs: list[dict[str, Any]],
    ) -> torch.Tensor:
        """Compute the task-specific embeddings.

        Args:
            features: The encoder features.
            inputs: The inputs to the model.

        Returns:
            The task-specific embeddings.
        """
        raise NotImplementedError

register_tasks 

register_tasks(task_names: list[str]) -> None

Register the tasks.

This must happen post-init so that we can dynamically determine the tasks to use, so it doesn't have to be specified in the config.

Parameters:

Name Type Description Default
task_names list[str]

The names of the tasks.

 required
Source code in rslearn/models/task_embedding.py 
55
56
57
58
59
60
61
62
63
64
def register_tasks(self, task_names: list[str]) -> None:
    """Register the tasks.

    This must happen post-init so that we can dynamically determine
    the tasks to use, so it doesn't have to be specified in the config.

    Args:
        task_names: The names of the tasks.
    """
    raise NotImplementedError

compute_embeds 

compute_embeds(features: list[tensor], inputs: list[dict[str, Any]]) -> Tensor

Compute the task-specific embeddings.

Parameters:

Name Type Description Default
features list[tensor]

The encoder features.

 required
inputs list[dict[str, Any]]

The inputs to the model.

 required

Returns:

Type Description
Tensor

The task-specific embeddings.
Source code in rslearn/models/task_embedding.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def compute_embeds(
    self,
    features: list[torch.tensor],
    inputs: list[dict[str, Any]],
) -> torch.Tensor:
    """Compute the task-specific embeddings.

    Args:
        features: The encoder features.
        inputs: The inputs to the model.

    Returns:
        The task-specific embeddings.
    """
    raise NotImplementedError

TaskChannelEmbedding

Bases: BaseTaskEmbedding

Registers task-specific 'tokens', i.e. embeddings.

Each embedding is learned per-channel and copied over the full spatial dimensions. Optionally, add a spatial sinusoidal positional embedding to the task embedding.

Source code in rslearn/models/task_embedding.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
class TaskChannelEmbedding(BaseTaskEmbedding):
    """Registers task-specific 'tokens', i.e. embeddings.

    Each embedding is learned per-channel and copied over the full spatial dimensions.
    Optionally, add a spatial sinusoidal positional embedding to the task embedding.
    """

    def __init__(
        self,
        encoder_embedding_size: int,
        default_idx: int = 0,
        add_spatial_embed: bool = False,
    ) -> None:
        """Initialize the task channel embedding module.

        Args:
            encoder_embedding_size: The size of the encoder embedding.
            default_idx: The index of the default task, useful if loading a merged model.
            add_spatial_embed: if true, add a spatial sinusoidal positional embedding to the task embedding
        """
        super().__init__(encoder_embedding_size)
        self.default_idx = default_idx
        self.add_spatial_embed = add_spatial_embed
        if add_spatial_embed:
            self.pos_embed = PositionalEncoding(encoder_embedding_size)

    def register_tasks(self, task_names: list[str]) -> None:
        """Register the tasks.

        This must happen post-init so that we can dynamically determine
        the tasks to use, so it doesn't have to be specified in the config.

        Args:
            task_names: The names of the tasks.
        """
        self.embed = torch.nn.Embedding(len(task_names), self.encoder_embedding_size)
        self.target_to_embed_idx = {name: i for i, name in enumerate(task_names)}

    def compute_embeds(
        self,
        features: list[torch.tensor],
        inputs: list[dict[str, Any]],
    ) -> torch.Tensor:
        """Compute the task-specific embeddings.

        Args:
            inputs: The inputs to the model.
            features: computed encoder features

        Returns:
            The task-specific embeddings, shape (B, T, C), T = HW
            The embeddings are repeated over the spatial dimensions, and optionally
            a sinusoidal positional embedding is added.
        """
        try:
            idx = [self.target_to_embed_idx[inp["dataset_source"]] for inp in inputs]
        except KeyError:
            idx = [self.default_idx] * len(inputs)
        embeds = self.embed(torch.tensor(idx).to(features[0].device))
        seq_len = features[0].shape[-1] * features[0].shape[-2]  # T = HW
        embeds = embeds.unsqueeze(0).repeat(seq_len, 1, 1)  # T x B x C
        if self.add_spatial_embed:
            embeds = self.pos_embed(embeds)
        embeds = torch.einsum("tbc->btc", embeds)  # B x T x C
        return embeds

    def forward(
        self,
        features: list[torch.tensor],
        inputs: list[dict[str, Any]],
        embeds: torch.Tensor | None = None,
    ) -> list[torch.tensor]:
        """Compute and apply task-specific embeddings to encoder features.

        Optionally, add a spatial sinusoidal positional embedding to the task embedding.
        Otherwise, the task embedding is repeated over the spatial dimensions.

        Args:
            features: The encoder features, a 1-list of B x C x H x W features.
            inputs: The inputs to the model.
            embeds: Already-computed task embeddings, if provided, skip the computation.

        Returns:
            The encoder features with the task-specific embeddings added.
        """
        height, width = features[0].shape[-2:]
        assert all(f.shape[-2:] == (height, width) for f in features), (
            "features must have the same spatial dimensions"
        )
        if embeds is None:
            embeds = self.compute_embeds(features, inputs)  # B x HW x C
        embeds = embeds.unflatten(dim=1, sizes=(height, width))  # B x H x W x C
        for i in range(len(features)):
            features[i] += torch.einsum("bhwc->bchw", embeds)  # B x C x H x W
        return features

register_tasks 

register_tasks(task_names: list[str]) -> None

Register the tasks.

This must happen post-init so that we can dynamically determine the tasks to use, so it doesn't have to be specified in the config.

Parameters:

Name Type Description Default
task_names list[str]

The names of the tasks.

 required
Source code in rslearn/models/task_embedding.py 
109
110
111
112
113
114
115
116
117
118
119
def register_tasks(self, task_names: list[str]) -> None:
    """Register the tasks.

    This must happen post-init so that we can dynamically determine
    the tasks to use, so it doesn't have to be specified in the config.

    Args:
        task_names: The names of the tasks.
    """
    self.embed = torch.nn.Embedding(len(task_names), self.encoder_embedding_size)
    self.target_to_embed_idx = {name: i for i, name in enumerate(task_names)}

compute_embeds 

compute_embeds(features: list[tensor], inputs: list[dict[str, Any]]) -> Tensor

Compute the task-specific embeddings.

Parameters:

Name Type Description Default
inputs list[dict[str, Any]]

The inputs to the model.

 required
features list[tensor]

computed encoder features

 required

Returns:

Type Description
Tensor

The task-specific embeddings, shape (B, T, C), T = HW
Tensor

The embeddings are repeated over the spatial dimensions, and optionally
Tensor

a sinusoidal positional embedding is added.
Source code in rslearn/models/task_embedding.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def compute_embeds(
    self,
    features: list[torch.tensor],
    inputs: list[dict[str, Any]],
) -> torch.Tensor:
    """Compute the task-specific embeddings.

    Args:
        inputs: The inputs to the model.
        features: computed encoder features

    Returns:
        The task-specific embeddings, shape (B, T, C), T = HW
        The embeddings are repeated over the spatial dimensions, and optionally
        a sinusoidal positional embedding is added.
    """
    try:
        idx = [self.target_to_embed_idx[inp["dataset_source"]] for inp in inputs]
    except KeyError:
        idx = [self.default_idx] * len(inputs)
    embeds = self.embed(torch.tensor(idx).to(features[0].device))
    seq_len = features[0].shape[-1] * features[0].shape[-2]  # T = HW
    embeds = embeds.unsqueeze(0).repeat(seq_len, 1, 1)  # T x B x C
    if self.add_spatial_embed:
        embeds = self.pos_embed(embeds)
    embeds = torch.einsum("tbc->btc", embeds)  # B x T x C
    return embeds

forward 

forward(features: list[tensor], inputs: list[dict[str, Any]], embeds: Tensor | None = None) -> list[tensor]

Compute and apply task-specific embeddings to encoder features.

Optionally, add a spatial sinusoidal positional embedding to the task embedding. Otherwise, the task embedding is repeated over the spatial dimensions.

Parameters:

Name Type Description Default
features list[tensor]

The encoder features, a 1-list of B x C x H x W features.

 required
inputs list[dict[str, Any]]

The inputs to the model.

 required
embeds Tensor | None

Already-computed task embeddings, if provided, skip the computation.

 None

Returns:

Type Description
list[tensor]

The encoder features with the task-specific embeddings added.
Source code in rslearn/models/task_embedding.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
def forward(
    self,
    features: list[torch.tensor],
    inputs: list[dict[str, Any]],
    embeds: torch.Tensor | None = None,
) -> list[torch.tensor]:
    """Compute and apply task-specific embeddings to encoder features.

    Optionally, add a spatial sinusoidal positional embedding to the task embedding.
    Otherwise, the task embedding is repeated over the spatial dimensions.

    Args:
        features: The encoder features, a 1-list of B x C x H x W features.
        inputs: The inputs to the model.
        embeds: Already-computed task embeddings, if provided, skip the computation.

    Returns:
        The encoder features with the task-specific embeddings added.
    """
    height, width = features[0].shape[-2:]
    assert all(f.shape[-2:] == (height, width) for f in features), (
        "features must have the same spatial dimensions"
    )
    if embeds is None:
        embeds = self.compute_embeds(features, inputs)  # B x HW x C
    embeds = embeds.unflatten(dim=1, sizes=(height, width))  # B x H x W x C
    for i in range(len(features)):
        features[i] += torch.einsum("bhwc->bchw", embeds)  # B x C x H x W
    return features

TaskMHAEmbedding

Bases: TaskChannelEmbedding

Multi-headed cross-attention over the spatial dimensions.

The task embedding is the query and the features are the key and value. We copy the task embedding over the spatial dimensions, and optionally add a sinusoidal positional embedding before the MHA layer.

Source code in rslearn/models/task_embedding.py 
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
class TaskMHAEmbedding(TaskChannelEmbedding):
    """Multi-headed cross-attention over the spatial dimensions.

    The task embedding is the query and the features are the key and value.
    We copy the task embedding over the spatial dimensions, and optionally
    add a sinusoidal positional embedding before the MHA layer.
    """

    def __init__(
        self,
        encoder_embedding_size: int,
        num_heads: int,
        default_idx: int = 0,
        add_spatial_embed: bool = True,
    ) -> None:
        """Initialize the task MHA embedding module.

        Args:
            encoder_embedding_size: The size of the encoder embedding.
            num_heads: The number of attention heads.
            default_idx: The index of the default task, useful if loading a merged model.
            add_spatial_embed: if true, add a spatial sinusoidal positional embedding to the task embedding
        """
        super().__init__(encoder_embedding_size, default_idx, add_spatial_embed)
        self.mha = torch.nn.MultiheadAttention(
            encoder_embedding_size, num_heads, batch_first=True
        )

    def register_tasks(self, task_names: list[str]) -> None:
        """Register the tasks.

        This must happen post-init so that we can dynamically determine
        the tasks to use, so it doesn't have to be specified in the config.

        Args:
            task_names: The names of the tasks.
        """
        super().register_tasks(task_names)

    def forward(
        self,
        features: list[torch.tensor],
        inputs: list[dict[str, Any]],
        embeds: torch.Tensor | None = None,
    ) -> list[torch.tensor]:
        """Compute and apply task-specific embeddings to encoder features.

        Also apply the MHA layer across the spatial dimension, with the task embedding
        as the query and the features as the key and value.

        Args:
            features: The encoder features, a 1-list of B x C x H x W features.
            inputs: The inputs to the model.
            embeds: Already-computed task embeddings, if provided, skip the computation.

        Returns:
            The encoder features with the task-specific embeddings added.
        """
        assert len(features) == 1, "TaskMHAEmbedding only supports one feature"
        x = torch.flatten(features[0], start_dim=2)  # B x C x T, T = HW
        if embeds is None:
            embeds = self.compute_embeds(features, inputs)  # B x T x C
        out = self.mha(
            embeds,  # B x T x C
            torch.einsum("bct->btc", x),
            torch.einsum("bct->btc", x),
        )[0]  # B x T x C
        out = torch.einsum("btc->bct", out)
        out = out.view(*features[0].shape)  # B x C x H x W
        return [out]

register_tasks 

register_tasks(task_names: list[str]) -> None

Register the tasks.

This must happen post-init so that we can dynamically determine the tasks to use, so it doesn't have to be specified in the config.

Parameters:

Name Type Description Default
task_names list[str]

The names of the tasks.

 required
Source code in rslearn/models/task_embedding.py 
208
209
210
211
212
213
214
215
216
217
def register_tasks(self, task_names: list[str]) -> None:
    """Register the tasks.

    This must happen post-init so that we can dynamically determine
    the tasks to use, so it doesn't have to be specified in the config.

    Args:
        task_names: The names of the tasks.
    """
    super().register_tasks(task_names)

forward 

forward(features: list[tensor], inputs: list[dict[str, Any]], embeds: Tensor | None = None) -> list[tensor]

Compute and apply task-specific embeddings to encoder features.

Also apply the MHA layer across the spatial dimension, with the task embedding as the query and the features as the key and value.

Parameters:

Name Type Description Default
features list[tensor]

The encoder features, a 1-list of B x C x H x W features.

 required
inputs list[dict[str, Any]]

The inputs to the model.

 required
embeds Tensor | None

Already-computed task embeddings, if provided, skip the computation.

 None

Returns:

Type Description
list[tensor]

The encoder features with the task-specific embeddings added.
Source code in rslearn/models/task_embedding.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
def forward(
    self,
    features: list[torch.tensor],
    inputs: list[dict[str, Any]],
    embeds: torch.Tensor | None = None,
) -> list[torch.tensor]:
    """Compute and apply task-specific embeddings to encoder features.

    Also apply the MHA layer across the spatial dimension, with the task embedding
    as the query and the features as the key and value.

    Args:
        features: The encoder features, a 1-list of B x C x H x W features.
        inputs: The inputs to the model.
        embeds: Already-computed task embeddings, if provided, skip the computation.

    Returns:
        The encoder features with the task-specific embeddings added.
    """
    assert len(features) == 1, "TaskMHAEmbedding only supports one feature"
    x = torch.flatten(features[0], start_dim=2)  # B x C x T, T = HW
    if embeds is None:
        embeds = self.compute_embeds(features, inputs)  # B x T x C
    out = self.mha(
        embeds,  # B x T x C
        torch.einsum("bct->btc", x),
        torch.einsum("bct->btc", x),
    )[0]  # B x T x C
    out = torch.einsum("btc->bct", out)
    out = out.view(*features[0].shape)  # B x C x H x W
    return [out]