[Feature Request] Support Multi-task learning #789

classicsong · 2024-03-29T05:18:05Z

GraphStorm should support multi-task learning as its built-in capability.

Unsupervised multi-task learning: Other than using link prediction as the sole supervision signal for unsupervised learning, one can add other supervision signals like reconstruct node features, predicting the motifs of specific nodes, etc (https://arxiv.org/abs/2007.10445).
Entities with different node types or edges from different edge types of a graph can have different tasks associated with them. There are also cases when the same node types or edge types have multiple tasks associated with them.

*Issue #, if available:* #789 *Description of changes:* This is the first PR to implement #789. We need to allow use to define train, validation and test mask names themselves. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

*Issue #, if available:* One PR of #789 Also solve #651 *Description of changes:* By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

…nk prediction dataloders reusable. (#825) *Issue #, if available:* Related to #789 *Description of changes:* 1. decouple the code of creating decoders from the code of creating built-in models to make the code of creating decoders more reusable. 2. add `get_builtin_lp_train_dataloader_class` and `get_builtin_lp_eval_dataloader_class` in gsf.py By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

…-task learning (#828) *Issue #, if available:* Support multi-task learning. First PR for #789 *Description of changes:* Update GraphStorm input config parsing to support multi-task learning. Allow user to specify to specify multiple training tasks for a training job through yaml file. By providing the `multi_task_learning` configurations in the yaml file, users can define multiple training tasks. The following config defines two training tasks, one for node classification and one for edge classification. ``` --- version: 1.0 gsf: basic: ... ... multi_task_learning: - node_classification: target_ntype: "movie" label_field: "label" mask_fields: - "train_mask_field_nc" - "val_mask_field_nc" - "test_mask_field_nc" task_weight: 1.0 - edge_classification: target_etype: - "user,rating,movie" label_field: "rate" mask_fields: - "train_mask_field_ec" - "val_mask_field_ec" - "test_mask_field_ec" task_weight: 0.5 # weight of the task ``` Task specific hyperparameters in multi-task learning are same as thoses in single task learning, except that two new configs are required, i.e., mask_fields and task_weight. The mask_fields provides the training, validation and test masks for the task and the task_weight gives its loss weight. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

…ng (#834) *Issue #, if available:* #789 *Description of changes:* Add GSgnnMultiTaskDataLoader to support multi-task learning. When initializing a GSgnnMultiTaskDataLoader, users need to provide two inputs: 1) a list of config.TaskInfo objects recording the information of each task and 2) a list of dataloaders corresponding to each training task. During training for each iteration, GSgnnMultiTaskDataLoader will iteratively call each task-dataloader to generate a mini-batch and finally return a list of mini-batches to the trainer. The length of the dataloader (number of batches for an epoch) is determined by the largest task in the GSgnnMultiTaskDataLoader. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

#837) *Issue #, if available:* #789 *Description of changes:* Add GSgnnMultiTaskEvaluator to support multi-task evaluation. GSgnnMultiTaskEvaluator accepts a set of Evaluators, in the format of dict ({task_id: Evaluator, ...}) as input to initialize the multi-task evaluator. When doing evaluation, it accepts three arguements val_results, test_results and total_iters. The val_results and test_results will be dicts in the format of {task_id_0: reslut, task_id_1: result}. The GSgnnMultiTaskEvaluator will call task specify evaluators for each task to compute the evaluation scores. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

*Issue #, if available:* #789 *Description of changes:* Add multi-task learning test dataset based on movielens. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

*Issue #, if available:* One PR of awslabs#789 Also solve awslabs#651 *Description of changes:* By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

…nk prediction dataloders reusable. (awslabs#825) *Issue #, if available:* Related to awslabs#789 *Description of changes:* 1. decouple the code of creating decoders from the code of creating built-in models to make the code of creating decoders more reusable. 2. add `get_builtin_lp_train_dataloader_class` and `get_builtin_lp_eval_dataloader_class` in gsf.py By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

…ng. (#852) *Issue #, if available:* #789 *Description of changes:* As multi-task learning trainer will invoke edge_mini_batch_predict, lp_mini_batch_predict and node_mini_batch_predict when conducting evaluation or testing, refactor the code to allow the functions to work with different decoders. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

*Issue #, if available:* #789 *Description of changes:* GSgnnMultiTaskSharedEncoderModel allows multiple tasks to share the same GNN encoder but have separate decoders for each task. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

*Issue #, if available:* #789 *Description of changes:* Add support for multi-task learning. Users can define multiple tasks in the same training loop. A task can be a node classification, node regression, edge classification, edge regression or link prediction task. For each node classification or node regression task, it should be defined on a single node type with one label field. But users can define multiple node classification or regression tasks on the same node type. For each edge classification or node regression task, it should be defined on a single edge type with one label field. But users can define multiple edge classification or regression tasks on the same edge type. For link prediction, users can define prediction targets on multiple edge types. ### Graph construction Update GraphStorm input config parsing to support multi-task learning. Allow user to specify multiple training tasks for a training job through yaml file. By providing the `multi_task_learning` configurations in the yaml file, users can define multiple training tasks. The following config defines two training tasks, one for node classification and one for edge classification. ``` --- version: 1.0 gsf: basic: ... ... multi_task_learning: - node_classification: target_ntype: "movie" label_field: "label" mask_fields: - "train_mask_field_nc" - "val_mask_field_nc" - "test_mask_field_nc" task_weight: 1.0 - edge_classification: target_etype: - "user,rating,movie" label_field: "rate" mask_fields: - "train_mask_field_ec" - "val_mask_field_ec" - "test_mask_field_ec" task_weight: 0.5 # weight of the task ``` Task specific hyperparameters in multi-task learning are same as thoses in single task learning, except that two new configs are required, i.e., mask_fields and task_weight. The mask_fields provides the training, validation and test masks for the task and the task_weight gives its loss weight. ### DataLoader for multi-task learning Add GSgnnMultiTaskDataLoader to support multi-task learning. When initializing a GSgnnMultiTaskDataLoader, users need to provide two inputs: 1) a list of config.TaskInfo objects recording the information of each task and 2) a list of dataloaders corresponding to each training task. During training for each iteration, GSgnnMultiTaskDataLoader will iteratively call each task-dataloader to generate a mini-batch and finally return a list of mini-batches to the trainer. The length of the dataloader (number of batches for an epoch) is determined by the largest task in the GSgnnMultiTaskDataLoader. #834 ### Evaluator for multi-task learning GSgnnMultiTaskEvaluator accepts a set of Evaluators, in the format of dict ({task_id: Evaluator, ...}) as input to initialize the multi-task evaluator. When doing evaluation, it accepts three arguements val_results, test_results and total_iters. The val_results and test_results will be dicts in the format of {task_id_0: reslut, task_id_1: result}. The GSgnnMultiTaskEvaluator will call task specify evaluators for each task to compute the evaluation scores. #837 ### Refactor graphstorm.model for multi-task learning As multi-task learning trainer will invoke edge_mini_batch_predict, lp_mini_batch_predict and node_mini_batch_predict when conducting evaluation or testing, refactor the code to allow the functions to work with different decoders. #843 ### Add GSgnnMultiTaskSharedEncoderModel GSgnnMultiTaskSharedEncoderModel allows multiple tasks to share the same GNN encoder but have separate decoders for each task. #855 ### Add Multi-task entrypoint #849 By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

) *Issue #, if available:* #789 #862 *Description of changes:* Add a new task type `BUILTIN_TASK_RECONSTRUCT_NODE_FEAT = "reconstruct_node_feat"`. User can use node feature reconstruction to supervise model training in multi-task learning. User can define a reconstruct_node_feat task as following: ``` ... multi_task_learning: - reconstruct_node_feat: reconstruct_nfeat_name: "title" target_ntype: "movie" batch_size: 128 mask_fields: - "train_mask_c0" # node classification mask 0 - "val_mask_c0" - "test_mask_c0" task_weight: 1.0 eval_metric: - "mse" ``` `reconstruct_node_feat` is the task name, `target_ntype` defines which node type, the reconstruct node feature learning will be applied. `reconstruct_nfeat_name` defines the name of the feature to be re-construct. The other configs are same as node regression tasks. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

*Issue #, if available:* #789 *Description of changes:* By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

… in lp and nfeat reconstruct task evaluation. (#871) *Issue #, if available:* #789 *Description of changes:* Previously, in the eval() function of GSgnnMultiTaskLearningTrainer, both link prediction and node feature reconstruction tasks use the node embeddings computed with the entire graph. This will cause test edge leakage for link prediction tasks and target node node feature leakage for node feature reconstruction tasks. This PR fixes this issue. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

*Issue #, if available:* #789 *Description of changes:* This PR adds inference support for multi-task learning. Users can use `python3 -m graphstorm.run.gs_multi_task_learning --inference ` to launch a inference task. This PR also changes remap_result.py to support remapping prediction results from multi-task learning inference. (The prediction results of each task are stored separately on different folders with the name of the corresponding task id.) By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

*Issue #, if available:* #789 *Description of changes:* Add doc for multi-task learning. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --------- Co-authored-by: Xiang Song <xiangsx@amazon.com>

classicsong mentioned this issue Mar 29, 2024

[Roadmap] 0.3 Release Plan. #788

Closed

15 tasks

classicsong mentioned this issue Apr 18, 2024

[GConstruct] Allow users define train/validation/test masks #804

Merged

classicsong mentioned this issue Apr 25, 2024

Support construct a graph for multi-task learning #813

Merged

classicsong mentioned this issue May 6, 2024

Refactor gsf.py to make the code of creating decoders and creating link prediction dataloders reusable. #825

Merged

classicsong mentioned this issue May 10, 2024

[Multi-task Learning] Update GraphStorm input config to support multi-task learning #828

Merged

This was referenced May 15, 2024

[Multi-task Learning] Add support for multi-task learning #842

Merged

[Multi-task Learning] Add multi-task evaluator for multi-task learning #837

Merged

[Multi-task Learning] Add multi-task Dataloader for multi-task learning #834

Merged

classicsong mentioned this issue May 15, 2024

[Multi-task Learning] Add Multi-task test dataset #843

Merged

classicsong mentioned this issue May 20, 2024

[Multi-task Learning] Refactor graphstorm.model for multi-task learning. #852

Merged

classicsong mentioned this issue May 23, 2024

[Multi-task Learning] Add GSgnnMultiTaskSharedEncoderModel #855

Merged

classicsong mentioned this issue May 30, 2024

Update gsgnn_node_emb for multi-task learning #860

Merged

This was referenced Jun 2, 2024

Support multi-task inference. #861

Merged

[Multi-task learning] Support reconstruct node feature supervision #863

Merged

classicsong mentioned this issue Jun 11, 2024

Update multi-task evaluation logic to avoid information leakage issue in lp and nfeat reconstruct task evaluation. #871

Merged

classicsong closed this as completed Jun 18, 2024

classicsong mentioned this issue Jun 20, 2024

[Documentation] Add doc for multi-task learning #887

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feature Request] Support Multi-task learning #789

[Feature Request] Support Multi-task learning #789

classicsong commented Mar 29, 2024

[Feature Request] Support Multi-task learning #789

[Feature Request] Support Multi-task learning #789

Comments

classicsong commented Mar 29, 2024