Skip to content

rslearn.train.tasks.multi_task

multi_task

Task for wrapping multiple tasks.

MultiTask

Bases: Task

A task for training on multiple tasks.

Source code in rslearn/train/tasks/multi_task.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
 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
class MultiTask(Task):
    """A task for training on multiple tasks."""

    def __init__(
        self, tasks: dict[str, Task], input_mapping: dict[str, dict[str, str]]
    ):
        """Create a new MultiTask.

        Args:
            tasks: map from task name to the task object
            input_mapping: for each task, maps which keys from the raw inputs should
                appear as potentially different keys for that task
        """
        self.tasks = tasks
        self.input_mapping = input_mapping

    def process_inputs(
        self,
        raw_inputs: dict[str, RasterImage | list[Feature]],
        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
        """
        input_dict = {}
        target_dict = {}
        if metadata.dataset_source is None:
            # No multi-dataset, so always compute across all tasks
            task_iter = list(self.tasks.items())
        else:
            # Multi-dataset, so only compute for the task in this dataset
            task_iter = [(metadata.dataset_source, self.tasks[metadata.dataset_source])]

        for task_name, task in task_iter:
            cur_raw_inputs = {}
            for k, v in self.input_mapping[task_name].items():
                if k not in raw_inputs:
                    continue
                cur_raw_inputs[v] = raw_inputs[k]

            cur_input_dict, cur_target_dict = task.process_inputs(
                cur_raw_inputs, metadata=metadata, load_targets=load_targets
            )
            input_dict[task_name] = cur_input_dict
            target_dict[task_name] = cur_target_dict

        return input_dict, target_dict

    def process_output(
        self, raw_output: Any, metadata: SampleMetadata
    ) -> dict[str, Any]:
        """Processes an output into raster or vector data.

        Args:
            raw_output: the output from prediction head. It must be a dict mapping from
                task name to per-task output for this sample.
            metadata: metadata about the patch being read

        Returns:
            either raster or vector data.
        """
        processed_output = {}
        for task_name, task in self.tasks.items():
            if task_name in raw_output:
                # In multi-dataset training, we may not have all datasets in the batch
                processed_output[task_name] = task.process_output(
                    raw_output[task_name], metadata
                )
        return processed_output

    def visualize(
        self,
        input_dict: dict[str, Any],
        target_dict: dict[str, Any] | None,
        output: dict[str, 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
        """
        images = {}
        for task_name, task in self.tasks.items():
            cur_target_dict = None
            if target_dict:
                cur_target_dict = target_dict[task_name]
            cur_images = task.visualize(input_dict, cur_target_dict, output[task_name])
            for label, image in cur_images.items():
                images[f"{task_name}_{label}"] = image
        return images

    def get_metrics(self) -> MetricCollection:
        """Get metrics for this task."""
        # Flatten metrics into a single dict with task_name/ prefix to avoid nested
        # MetricCollections. Nested collections cause issues because MetricCollection
        # has postfix=None which breaks MetricCollection.compute().
        all_metrics = {}
        for task_name, task in self.tasks.items():
            for metric_name, metric in task.get_metrics().items():
                all_metrics[f"{task_name}/{metric_name}"] = MetricWrapper(
                    task_name, metric
                )
        return MetricCollection(all_metrics)

process_inputs 

process_inputs(raw_inputs: dict[str, RasterImage | list[Feature]], 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, RasterImage | list[Feature]]

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/multi_task.py 
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
def process_inputs(
    self,
    raw_inputs: dict[str, RasterImage | list[Feature]],
    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
    """
    input_dict = {}
    target_dict = {}
    if metadata.dataset_source is None:
        # No multi-dataset, so always compute across all tasks
        task_iter = list(self.tasks.items())
    else:
        # Multi-dataset, so only compute for the task in this dataset
        task_iter = [(metadata.dataset_source, self.tasks[metadata.dataset_source])]

    for task_name, task in task_iter:
        cur_raw_inputs = {}
        for k, v in self.input_mapping[task_name].items():
            if k not in raw_inputs:
                continue
            cur_raw_inputs[v] = raw_inputs[k]

        cur_input_dict, cur_target_dict = task.process_inputs(
            cur_raw_inputs, metadata=metadata, load_targets=load_targets
        )
        input_dict[task_name] = cur_input_dict
        target_dict[task_name] = cur_target_dict

    return input_dict, target_dict

process_output 

process_output(raw_output: Any, metadata: SampleMetadata) -> dict[str, Any]

Processes an output into raster or vector data.

Parameters:

Name Type Description Default
raw_output Any

the output from prediction head. It must be a dict mapping from task name to per-task output for this sample.

 required
metadata SampleMetadata

metadata about the patch being read

 required

Returns:

Type Description
dict[str, Any]

either raster or vector data.
Source code in rslearn/train/tasks/multi_task.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def process_output(
    self, raw_output: Any, metadata: SampleMetadata
) -> dict[str, Any]:
    """Processes an output into raster or vector data.

    Args:
        raw_output: the output from prediction head. It must be a dict mapping from
            task name to per-task output for this sample.
        metadata: metadata about the patch being read

    Returns:
        either raster or vector data.
    """
    processed_output = {}
    for task_name, task in self.tasks.items():
        if task_name in raw_output:
            # In multi-dataset training, we may not have all datasets in the batch
            processed_output[task_name] = task.process_output(
                raw_output[task_name], metadata
            )
    return processed_output

visualize 

visualize(input_dict: dict[str, Any], target_dict: dict[str, Any] | None, output: dict[str, 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 dict[str, 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/multi_task.py 
 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
def visualize(
    self,
    input_dict: dict[str, Any],
    target_dict: dict[str, Any] | None,
    output: dict[str, 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
    """
    images = {}
    for task_name, task in self.tasks.items():
        cur_target_dict = None
        if target_dict:
            cur_target_dict = target_dict[task_name]
        cur_images = task.visualize(input_dict, cur_target_dict, output[task_name])
        for label, image in cur_images.items():
            images[f"{task_name}_{label}"] = image
    return images

get_metrics 

get_metrics() -> MetricCollection

Get metrics for this task.

Source code in rslearn/train/tasks/multi_task.py 
119
120
121
122
123
124
125
126
127
128
129
130
def get_metrics(self) -> MetricCollection:
    """Get metrics for this task."""
    # Flatten metrics into a single dict with task_name/ prefix to avoid nested
    # MetricCollections. Nested collections cause issues because MetricCollection
    # has postfix=None which breaks MetricCollection.compute().
    all_metrics = {}
    for task_name, task in self.tasks.items():
        for metric_name, metric in task.get_metrics().items():
            all_metrics[f"{task_name}/{metric_name}"] = MetricWrapper(
                task_name, metric
            )
    return MetricCollection(all_metrics)

MetricWrapper

Bases: Metric

Wrapper for a metric from one task to operate in the multi-task setting.

It selects the outputs and targets that are relevant to each task.

Source code in rslearn/train/tasks/multi_task.py 
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
178
179
180
181
182
183
184
185
186
187
188
189
190
191
class MetricWrapper(Metric):
    """Wrapper for a metric from one task to operate in the multi-task setting.

    It selects the outputs and targets that are relevant to each task.
    """

    def __init__(self, task_name: str, metric: Metric):
        """Create a new MetricWrapper.

        The wrapper passes the task-specific predictions and targets to the metrics of
        returned from each task.

        Args:
            task_name: the name of the task
            metric: one metric from the task to wrap
        """
        super().__init__()
        self.task_name = task_name
        self.metric = metric

    def update(
        self, preds: list[dict[str, Any]], targets: list[dict[str, Any]]
    ) -> None:
        """Update metric.

        Args:
            preds: the predictions
            targets: the targets
        """
        try:
            self.metric.update(
                [pred[self.task_name] for pred in preds],
                [target[self.task_name] for target in targets],
            )
        except KeyError:
            # In multi-dataset training, we may not have all datasets in the batch
            pass

    def compute(self) -> Any:
        """Returns the computed metric."""
        result = self.metric.compute()
        # If the inner metric returns a dict, prefix each key with the task name so
        # keys don't clash across tasks in MetricCollection. For scalar metrics, the
        # name in the MetricCollection will be included as a prefix in the dict
        # returned from MetricCollection.compute(), but for dict metrics, it copies
        # over the key directly with no prefixing, so we need to handle the prefixing
        # ourselves.
        if isinstance(result, dict):
            return {f"{self.task_name}/{k}": v for k, v in result.items()}
        return result

    def reset(self) -> None:
        """Reset metric."""
        super().reset()
        self.metric.reset()

    def plot(self, *args: list[Any], **kwargs: dict[str, Any]) -> Any:
        """Returns a plot of the metric."""
        return self.metric.plot(*args, **kwargs)

update 

update(preds: list[dict[str, Any]], targets: list[dict[str, Any]]) -> None

Update metric.

Parameters:

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

the predictions

 required
targets list[dict[str, Any]]

the targets

 required
Source code in rslearn/train/tasks/multi_task.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def update(
    self, preds: list[dict[str, Any]], targets: list[dict[str, Any]]
) -> None:
    """Update metric.

    Args:
        preds: the predictions
        targets: the targets
    """
    try:
        self.metric.update(
            [pred[self.task_name] for pred in preds],
            [target[self.task_name] for target in targets],
        )
    except KeyError:
        # In multi-dataset training, we may not have all datasets in the batch
        pass

compute 

compute() -> Any

Returns the computed metric.

Source code in rslearn/train/tasks/multi_task.py 
171
172
173
174
175
176
177
178
179
180
181
182
def compute(self) -> Any:
    """Returns the computed metric."""
    result = self.metric.compute()
    # If the inner metric returns a dict, prefix each key with the task name so
    # keys don't clash across tasks in MetricCollection. For scalar metrics, the
    # name in the MetricCollection will be included as a prefix in the dict
    # returned from MetricCollection.compute(), but for dict metrics, it copies
    # over the key directly with no prefixing, so we need to handle the prefixing
    # ourselves.
    if isinstance(result, dict):
        return {f"{self.task_name}/{k}": v for k, v in result.items()}
    return result

reset 

reset() -> None

Reset metric.

Source code in rslearn/train/tasks/multi_task.py 
184
185
186
187
def reset(self) -> None:
    """Reset metric."""
    super().reset()
    self.metric.reset()

plot 

plot(*args: list[Any], **kwargs: dict[str, Any]) -> Any

Returns a plot of the metric.

Source code in rslearn/train/tasks/multi_task.py 
189
190
191
def plot(self, *args: list[Any], **kwargs: dict[str, Any]) -> Any:
    """Returns a plot of the metric."""
    return self.metric.plot(*args, **kwargs)