Skip to content

rslearn.train.tasks.embedding

embedding

Embedding task.

EmbeddingTask

Bases: Task

A dummy task for computing embeddings.

This task does not compute any targets or loss. Instead, it is just set up for inference, to save embeddings from the configured model.

Source code in rslearn/train/tasks/embedding.py 
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
82
83
class EmbeddingTask(Task):
    """A dummy task for computing embeddings.

    This task does not compute any targets or loss. Instead, it is just set up for
    inference, to save embeddings from the configured model.
    """

    def process_inputs(
        self,
        raw_inputs: dict[str, torch.Tensor],
        metadata: SampleMetadata,
        load_targets: bool = True,
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """Processes the data into targets.

        Args:
            raw_inputs: raster or vector data to process
            metadata: metadata about the patch being read
            load_targets: whether to load the targets or only inputs

        Returns:
            tuple (input_dict, target_dict) containing the processed inputs and targets
                that are compatible with both metrics and loss functions
        """
        return {}, {}

    def process_output(
        self, raw_output: Any, metadata: SampleMetadata
    ) -> npt.NDArray[Any] | list[Feature]:
        """Processes an output into raster or vector data.

        Args:
            raw_output: the output from prediction head, which must be a CxHxW tensor.
            metadata: metadata about the patch being read

        Returns:
            either raster or vector data.
        """
        if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 3:
            raise ValueError(
                "output for EmbeddingTask must be a tensor with three dimensions"
            )

        # Just convert the raw output to numpy array that can be saved to GeoTIFF.
        return raw_output.cpu().numpy()

    def visualize(
        self,
        input_dict: dict[str, Any],
        target_dict: dict[str, Any] | None,
        output: Any,
    ) -> dict[str, npt.NDArray[Any]]:
        """Visualize the outputs and targets.

        Args:
            input_dict: the input dict from process_inputs
            target_dict: the target dict from process_inputs
            output: the prediction

        Returns:
            a dictionary mapping image name to visualization image
        """
        # EmbeddingTask is only set up to support `model predict`.
        raise NotImplementedError

    def get_metrics(self) -> MetricCollection:
        """Get the metrics for this task."""
        return MetricCollection({})

process_inputs 

process_inputs(raw_inputs: dict[str, Tensor], metadata: SampleMetadata, load_targets: bool = True) -> tuple[dict[str, Any], dict[str, Any]]

Processes the data into targets.

Parameters:

Name Type Description Default
raw_inputs dict[str, Tensor]

raster or vector data to process

 required
metadata SampleMetadata

metadata about the patch being read

 required
load_targets bool

whether to load the targets or only inputs

 True

Returns:

Type Description
tuple[dict[str, Any], dict[str, Any]]

tuple (input_dict, target_dict) containing the processed inputs and targets that are compatible with both metrics and loss functions
Source code in rslearn/train/tasks/embedding.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def process_inputs(
    self,
    raw_inputs: dict[str, torch.Tensor],
    metadata: SampleMetadata,
    load_targets: bool = True,
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Processes the data into targets.

    Args:
        raw_inputs: raster or vector data to process
        metadata: metadata about the patch being read
        load_targets: whether to load the targets or only inputs

    Returns:
        tuple (input_dict, target_dict) containing the processed inputs and targets
            that are compatible with both metrics and loss functions
    """
    return {}, {}

process_output 

process_output(raw_output: Any, metadata: SampleMetadata) -> NDArray[Any] | list[Feature]

Processes an output into raster or vector data.

Parameters:

Name Type Description Default
raw_output Any

the output from prediction head, which must be a CxHxW tensor.

 required
metadata SampleMetadata

metadata about the patch being read

 required

Returns:

Type Description
NDArray[Any] | list[Feature]

either raster or vector data.
Source code in rslearn/train/tasks/embedding.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def process_output(
    self, raw_output: Any, metadata: SampleMetadata
) -> npt.NDArray[Any] | list[Feature]:
    """Processes an output into raster or vector data.

    Args:
        raw_output: the output from prediction head, which must be a CxHxW tensor.
        metadata: metadata about the patch being read

    Returns:
        either raster or vector data.
    """
    if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 3:
        raise ValueError(
            "output for EmbeddingTask must be a tensor with three dimensions"
        )

    # Just convert the raw output to numpy array that can be saved to GeoTIFF.
    return raw_output.cpu().numpy()

visualize 

visualize(input_dict: dict[str, Any], target_dict: dict[str, Any] | None, output: Any) -> dict[str, NDArray[Any]]

Visualize the outputs and targets.

Parameters:

Name Type Description Default
input_dict dict[str, Any]

the input dict from process_inputs

 required
target_dict dict[str, Any] | None

the target dict from process_inputs

 required
output Any

the prediction

 required

Returns:

Type Description
dict[str, NDArray[Any]]

a dictionary mapping image name to visualization image
Source code in rslearn/train/tasks/embedding.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def visualize(
    self,
    input_dict: dict[str, Any],
    target_dict: dict[str, Any] | None,
    output: Any,
) -> dict[str, npt.NDArray[Any]]:
    """Visualize the outputs and targets.

    Args:
        input_dict: the input dict from process_inputs
        target_dict: the target dict from process_inputs
        output: the prediction

    Returns:
        a dictionary mapping image name to visualization image
    """
    # EmbeddingTask is only set up to support `model predict`.
    raise NotImplementedError

get_metrics 

get_metrics() -> MetricCollection

Get the metrics for this task.

Source code in rslearn/train/tasks/embedding.py 
81
82
83
def get_metrics(self) -> MetricCollection:
    """Get the metrics for this task."""
    return MetricCollection({})

EmbeddingHead

Bases: Predictor

Head for embedding task.

It just adds a dummy loss to act as a Predictor.

Source code in rslearn/train/tasks/embedding.py 
 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
class EmbeddingHead(Predictor):
    """Head for embedding task.

    It just adds a dummy loss to act as a Predictor.
    """

    def forward(
        self,
        intermediates: Any,
        context: ModelContext,
        targets: list[dict[str, Any]] | None = None,
    ) -> ModelOutput:
        """Return the feature map along with a dummy loss.

        Args:
            intermediates: output from the previous model component, which must be a
                FeatureMaps consisting of a single feature map.
            context: the model context.
            targets: the targets (ignored).

        Returns:
            model output with the feature map that was input to this component along
                with a dummy loss.
        """
        if not isinstance(intermediates, FeatureMaps):
            raise ValueError("input to EmbeddingHead must be a FeatureMaps")
        if len(intermediates.feature_maps) != 1:
            raise ValueError(
                f"input to EmbeddingHead must have one feature map, but got {len(intermediates.feature_maps)}"
            )

        return ModelOutput(
            outputs=intermediates.feature_maps[0],
            loss_dict={"loss": 0},
        )

forward 

forward(intermediates: Any, context: ModelContext, targets: list[dict[str, Any]] | None = None) -> ModelOutput

Return the feature map along with a dummy loss.

Parameters:

Name Type Description Default
intermediates Any

output from the previous model component, which must be a FeatureMaps consisting of a single feature map.

 required
context ModelContext

the model context.

 required
targets list[dict[str, Any]] | None

the targets (ignored).

 None

Returns:

Type Description
ModelOutput

model output with the feature map that was input to this component along with a dummy loss.
Source code in rslearn/train/tasks/embedding.py 
 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
def forward(
    self,
    intermediates: Any,
    context: ModelContext,
    targets: list[dict[str, Any]] | None = None,
) -> ModelOutput:
    """Return the feature map along with a dummy loss.

    Args:
        intermediates: output from the previous model component, which must be a
            FeatureMaps consisting of a single feature map.
        context: the model context.
        targets: the targets (ignored).

    Returns:
        model output with the feature map that was input to this component along
            with a dummy loss.
    """
    if not isinstance(intermediates, FeatureMaps):
        raise ValueError("input to EmbeddingHead must be a FeatureMaps")
    if len(intermediates.feature_maps) != 1:
        raise ValueError(
            f"input to EmbeddingHead must have one feature map, but got {len(intermediates.feature_maps)}"
        )

    return ModelOutput(
        outputs=intermediates.feature_maps[0],
        loss_dict={"loss": 0},
    )