Skip to content

rslearn.models.trunk

trunk

Trunk module for decoder.

DecoderTrunkLayer

Bases: Module, ABC

Trunk layer for decoder.

Source code in rslearn/models/trunk.py 
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
class DecoderTrunkLayer(torch.nn.Module, ABC):
    """Trunk layer for decoder."""

    def __init__(self) -> None:
        """Initialize the DecoderTrunkLayer module."""
        super().__init__()

    @abstractmethod
    def forward(
        self, x: torch.Tensor, task_embedding: torch.Tensor | None = None
    ) -> dict[str, torch.Tensor]:
        """Forward pass.

        Args:
            x: input tensor of shape (batch_size, seq_len, dim)
            task_embedding: task embedding tensor of shape (batch_size, dim), or None

        Returns:
            dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
            and optionally other keys.
        """
        raise NotImplementedError

    @abstractmethod
    def apply_auxiliary_losses(
        self, trunk_out: dict[str, Any], outs: ModelOutput
    ) -> None:
        """Apply auxiliary losses in-place.

        Args:
            trunk_out: The output of the trunk.
            outs: The output of the decoders, with key "loss_dict" containing the losses.
        """
        raise NotImplementedError

forward abstractmethod 

forward(x: Tensor, task_embedding: Tensor | None = None) -> dict[str, Tensor]

Forward pass.

Parameters:

Name Type Description Default
x Tensor

input tensor of shape (batch_size, seq_len, dim)

 required
task_embedding Tensor | None

task embedding tensor of shape (batch_size, dim), or None

 None

Returns:

Type Description
dict[str, Tensor]

dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
dict[str, Tensor]

and optionally other keys.
Source code in rslearn/models/trunk.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
@abstractmethod
def forward(
    self, x: torch.Tensor, task_embedding: torch.Tensor | None = None
) -> dict[str, torch.Tensor]:
    """Forward pass.

    Args:
        x: input tensor of shape (batch_size, seq_len, dim)
        task_embedding: task embedding tensor of shape (batch_size, dim), or None

    Returns:
        dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
        and optionally other keys.
    """
    raise NotImplementedError

apply_auxiliary_losses abstractmethod 

apply_auxiliary_losses(trunk_out: dict[str, Any], outs: ModelOutput) -> None

Apply auxiliary losses in-place.

Parameters:

Name Type Description Default
trunk_out dict[str, Any]

The output of the trunk.

 required
outs ModelOutput

The output of the decoders, with key "loss_dict" containing the losses.

 required
Source code in rslearn/models/trunk.py 
38
39
40
41
42
43
44
45
46
47
48
@abstractmethod
def apply_auxiliary_losses(
    self, trunk_out: dict[str, Any], outs: ModelOutput
) -> None:
    """Apply auxiliary losses in-place.

    Args:
        trunk_out: The output of the trunk.
        outs: The output of the decoders, with key "loss_dict" containing the losses.
    """
    raise NotImplementedError

DecoderTrunk

Bases: Module

Trunk module for decoder, including arbitrary layers plus an optional task embedding.

Source code in rslearn/models/trunk.py 
 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
 82
 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
class DecoderTrunk(torch.nn.Module):
    """Trunk module for decoder, including arbitrary layers plus an optional task embedding."""

    def __init__(
        self,
        task_embedding: BaseTaskEmbedding | None = None,
        layers: list[DecoderTrunkLayer] | None = None,
    ) -> None:
        """Initialize the DecoderTrunk module.

        Args:
            task_embedding: Task-specific embedding module, or None if not using task embedding.
            layers: List of other shared layers. The first one should expect a
                B x T x C tensor, and the last should output a B x T x C tensor.
                All layers must output a dict with key "outputs" (output tensor of shape
                (B, T, C)) and optionally other keys.
        """
        super().__init__()
        self.layers = torch.nn.ModuleList(layers or [])
        self.task_embedding = task_embedding

        # If we have multiple instances of the same layer class, output keys will get overwritten
        if layers is not None:
            types = [type(layer) for layer in layers]
            if len(set(types)) != len(types):
                logger.warning(
                    "Multiple instances of the same layer class found in trunk. "
                    "Only the keys from the last instance will be used"
                )

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

        Args:
            task_names: list of task names
        """
        if self.task_embedding is not None:
            self.task_embedding.register_tasks(task_names)

    def forward(
        self,
        features: list[torch.tensor],
        inputs: list[dict[str, Any]],
    ) -> dict[str, Any]:
        """Forward pass.

        Args:
            features: The encoder features, a 1-list of B x C x H x W features.
            inputs: The original inputs to the encoder.

        Returns:
            dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
            and optionally other keys from the other layers.
        """
        embeds = None
        if self.task_embedding is not None:
            embeds = self.task_embedding.compute_embeds(features, inputs)
            features = self.task_embedding(features, inputs, embeds=embeds)

        if not self.layers:
            return {"outputs": features}

        assert len(features) == 1, "DecoderTrunk only supports one feature map"
        x = torch.einsum("bchw->bhwc", features[0])
        x = torch.flatten(x, start_dim=1, end_dim=2)  # B x T x C, T = HW
        out = {}
        for layer in self.layers:
            layer_out = layer(x, task_embedding=embeds)
            x = layer_out.pop("outputs")  # unspecified shape
            out.update(layer_out)
        x = torch.einsum("btc->bct", x)  # B x C x T
        x = x.view(*features[0].shape)  # B x C x H x W

        out["outputs"] = [x]
        return out

    def apply_auxiliary_losses(
        self, trunk_out: dict[str, Any], outs: ModelOutput
    ) -> None:
        """Apply auxiliary losses in-place.

        Each layer handles its own auxiliary losses, assuming the loss key is `loss_dict`.

        Args:
            trunk_out: The output of the trunk.
            outs: The output of the decoders.
        """
        for layer in self.layers:
            layer.apply_auxiliary_losses(trunk_out, outs)

register_tasks 

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

Register tasks.

Parameters:

Name Type Description Default
task_names list[str]

list of task names

 required
Source code in rslearn/models/trunk.py 
81
82
83
84
85
86
87
88
def register_tasks(self, task_names: list[str]) -> None:
    """Register tasks.

    Args:
        task_names: list of task names
    """
    if self.task_embedding is not None:
        self.task_embedding.register_tasks(task_names)

forward 

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

Forward pass.

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 original inputs to the encoder.

 required

Returns:

Type Description
dict[str, Any]

dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
dict[str, Any]

and optionally other keys from the other layers.
Source code in rslearn/models/trunk.py 
 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
def forward(
    self,
    features: list[torch.tensor],
    inputs: list[dict[str, Any]],
) -> dict[str, Any]:
    """Forward pass.

    Args:
        features: The encoder features, a 1-list of B x C x H x W features.
        inputs: The original inputs to the encoder.

    Returns:
        dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
        and optionally other keys from the other layers.
    """
    embeds = None
    if self.task_embedding is not None:
        embeds = self.task_embedding.compute_embeds(features, inputs)
        features = self.task_embedding(features, inputs, embeds=embeds)

    if not self.layers:
        return {"outputs": features}

    assert len(features) == 1, "DecoderTrunk only supports one feature map"
    x = torch.einsum("bchw->bhwc", features[0])
    x = torch.flatten(x, start_dim=1, end_dim=2)  # B x T x C, T = HW
    out = {}
    for layer in self.layers:
        layer_out = layer(x, task_embedding=embeds)
        x = layer_out.pop("outputs")  # unspecified shape
        out.update(layer_out)
    x = torch.einsum("btc->bct", x)  # B x C x T
    x = x.view(*features[0].shape)  # B x C x H x W

    out["outputs"] = [x]
    return out

apply_auxiliary_losses 

apply_auxiliary_losses(trunk_out: dict[str, Any], outs: ModelOutput) -> None

Apply auxiliary losses in-place.

Each layer handles its own auxiliary losses, assuming the loss key is loss_dict.

Parameters:

Name Type Description Default
trunk_out dict[str, Any]

The output of the trunk.

 required
outs ModelOutput

The output of the decoders.

 required
Source code in rslearn/models/trunk.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
def apply_auxiliary_losses(
    self, trunk_out: dict[str, Any], outs: ModelOutput
) -> None:
    """Apply auxiliary losses in-place.

    Each layer handles its own auxiliary losses, assuming the loss key is `loss_dict`.

    Args:
        trunk_out: The output of the trunk.
        outs: The output of the decoders.
    """
    for layer in self.layers:
        layer.apply_auxiliary_losses(trunk_out, outs)