[Multi-task Learning] Add multi-task trainer

.github/workflow_scripts/e2e_mgpu_check.sh

@@ -10,4 +10,5 @@ sh ./tests/end2end-tests/create_data.sh
 bash ./tests/end2end-tests/graphstorm-lp/mgpu_test.sh
 bash ./tests/end2end-tests/graphstorm-nc/mgpu_test.sh
 bash ./tests/end2end-tests/graphstorm-ec/mgpu_test.sh
+bash ./tests/end2end-tests/graphstorm-mt/mgpu_test.sh
 

python/graphstorm/__init__.py

@@ -28,5 +28,10 @@
 from .gsf import create_builtin_lp_model
 from .gsf import create_builtin_edge_model
 from .gsf import create_builtin_node_model
+from .gsf import create_task_decoder
+
+from .gsf import (create_builtin_node_decoder,
+                  create_builtin_edge_decoder,
+                  create_builtin_lp_decoder)
 from .gsf import (get_builtin_lp_train_dataloader_class,
                   get_builtin_lp_eval_dataloader_class)

python/graphstorm/config/argument.py

@@ -150,9 +150,9 @@ def __init__(self, cmd_args):
         configuration = self.load_yaml_config(cmd_args.yaml_config_file)
 
         multi_task_config = None
-        if 'multi_task_learning' in configuration:
-            multi_task_config = configuration['multi_task_learning']
-            del configuration['multi_task_learning']
+        if 'multi_task_learning' in configuration['gsf']:
+            multi_task_config = configuration['gsf']['multi_task_learning']
+            del configuration['gsf']['multi_task_learning']
 
         self.set_attributes(configuration)
         # Override class attributes using command-line arguments
@@ -305,7 +305,9 @@ def _parse_node_classification_task(self, task_config):
         task_id = get_mttask_id(task_type=task_type,
                                 ntype=target_ntype,
                                 label=label_field)
-        setattr(task_info, "mask_fields", mask_fields)
+        setattr(task_info, "train_mask", mask_fields[0])
+        setattr(task_info, "val_mask", mask_fields[1])
+        setattr(task_info, "test_mask", mask_fields[2])
         setattr(task_info, "task_weight", task_weight)
 
         return TaskInfo(task_type=task_type,
@@ -336,7 +338,9 @@ def _parse_node_regression_task(self, task_config):
         task_id = get_mttask_id(task_type=task_type,
                                 ntype=target_ntype,
                                 label=label_field)
-        setattr(task_info, "mask_fields", mask_fields)
+        setattr(task_info, "train_mask", mask_fields[0])
+        setattr(task_info, "val_mask", mask_fields[1])
+        setattr(task_info, "test_mask", mask_fields[2])
         setattr(task_info, "task_weight", task_weight)
 
         return TaskInfo(task_type=task_type,
@@ -367,7 +371,9 @@ def _parse_edge_classification_task(self, task_config):
         task_id = get_mttask_id(task_type=task_type,
                                 etype=target_etype,
                                 label=label_field)
-        setattr(task_info, "mask_fields", mask_fields)
+        setattr(task_info, "train_mask", mask_fields[0])
+        setattr(task_info, "val_mask", mask_fields[1])
+        setattr(task_info, "test_mask", mask_fields[2])
         setattr(task_info, "task_weight", task_weight)
         return TaskInfo(task_type=task_type,
                         task_id=task_id,
@@ -397,7 +403,9 @@ def _parse_edge_regression_task(self, task_config):
         task_id = get_mttask_id(task_type=task_type,
                                 etype=target_etype,
                                 label=label_field)
-        setattr(task_info, "mask_fields", mask_fields)
+        setattr(task_info, "train_mask", mask_fields[0])
+        setattr(task_info, "val_mask", mask_fields[1])
+        setattr(task_info, "test_mask", mask_fields[2])
         setattr(task_info, "task_weight", task_weight)
         return TaskInfo(task_type=task_type,
                         task_id=task_id,
@@ -425,7 +433,9 @@ def _parse_link_prediction_task(self, task_config):
         task_id = get_mttask_id(
             task_type=task_type,
             etype=train_etype if train_etype is not None else "ALL_ETYPE")
-        setattr(task_info, "mask_fields", mask_fields)
+        setattr(task_info, "train_mask", mask_fields[0])
+        setattr(task_info, "val_mask", mask_fields[1])
+        setattr(task_info, "test_mask", mask_fields[2])
         setattr(task_info, "task_weight", task_weight)
         return TaskInfo(task_type=task_type,
                         task_id=task_id,

python/graphstorm/gsf.py

@@ -28,10 +28,11 @@
 
 from .utils import sys_tracker, get_rank
 from .utils import setup_device
-from .config import BUILTIN_TASK_NODE_CLASSIFICATION
-from .config import BUILTIN_TASK_NODE_REGRESSION
-from .config import BUILTIN_TASK_EDGE_CLASSIFICATION
-from .config import BUILTIN_TASK_EDGE_REGRESSION
+from graphstorm.config import (BUILTIN_TASK_NODE_CLASSIFICATION,
+                               BUILTIN_TASK_NODE_REGRESSION,
+                               BUILTIN_TASK_EDGE_CLASSIFICATION,
+                               BUILTIN_TASK_EDGE_REGRESSION,
+                               BUILTIN_TASK_LINK_PREDICTION)
 from .config import BUILTIN_LP_DOT_DECODER
 from .config import BUILTIN_LP_DISTMULT_DECODER
 from .config import (BUILTIN_LP_LOSS_CROSS_ENTROPY,
@@ -842,3 +843,31 @@ def get_builtin_lp_train_dataloader_class(config):
         raise ValueError('Unknown negative sampler')
 
     return dataloader_cls
+
+def create_task_decoder(task_info, g, decoder_input_dim, train_task):
+    """ Create task decoders according to task_info.
+
+    Parameters
+    ----------
+    task_info: TaskInfo
+        Task info.
+    g: Dist DGLGraph
+        Graph
+    decoder_input_dim: int
+        The dimension of the input embedding of the decoder
+    train_task: bool
+        Whether the task is a training task
+
+    Return
+    ------
+    decoder: The node task decoder(s)
+    loss_func: The loss function(s)
+    """
+    if task_info.task_type in [BUILTIN_TASK_NODE_CLASSIFICATION, BUILTIN_TASK_NODE_REGRESSION]:
+        return create_builtin_node_decoder(g, decoder_input_dim, task_info.task_config, train_task)
+    elif task_info.task_type in [BUILTIN_TASK_EDGE_CLASSIFICATION, BUILTIN_TASK_EDGE_REGRESSION]:
+        return create_builtin_edge_decoder(g, decoder_input_dim, task_info.task_config, train_task)
+    elif task_info.task_type in [BUILTIN_TASK_LINK_PREDICTION]:
+        return create_builtin_lp_decoder(g, decoder_input_dim, task_info.task_config, train_task)
+
+    return None, None

python/graphstorm/inference/lp_infer.py

@@ -72,7 +72,7 @@ def infer(self, data, loader, save_embed_path,
         save_embed_format : str
             Specify the format of saved embeddings.
         infer_batch_size: int
-            Specify the inference batch size when computing node embeddings 
+            Specify the inference batch size when computing node embeddings
             with mini batch inference.
         """
         sys_tracker.check('start inferencing')

python/graphstorm/run/gs_multi_task_learning.py

@@ -0,0 +1,52 @@
+"""
+    Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+    Licensed under the Apache License, Version 2.0 (the "License").
+    You may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+        http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+
+    Entry point for running multi-task learning.
+
+    Run as:
+    python3 -m graphstorm.run.gs_multi_task_learning <Launch args> <Train/Infer args>
+"""
+import os
+import logging
+
+from .launch import get_argument_parser
+from .launch import check_input_arguments
+from .launch import submit_jobs
+
+def main():
+    """ Main function
+    """
+    parser = get_argument_parser()
+    args, exec_script_args = parser.parse_known_args()
+    check_input_arguments(args)
+
+    lib_dir = os.path.abspath(os.path.dirname(__file__))
+    if args.inference:
+        cmd = "gsgnn_mt/gsgnn_infer_mt.py"
+    else:
+        cmd = "gsgnn_mt/gsgnn_mt.py"
+    cmd_path = os.path.join(lib_dir, cmd)
+    exec_script_args = [cmd_path] + exec_script_args
+
+    if "coo" not in args.graph_format:
+        args.graph_format = f"{args.graph_format},coo"
+        logging.debug("Automatically add COO format to graph formats for link prediction. " + \
+                "New graph_format is %s", args.graph_format)
+    submit_jobs(args, exec_script_args)
+
+if __name__ == "__main__":
+    FMT = "%(asctime)s %(levelname)s %(message)s"
+    logging.basicConfig(format=FMT, level=logging.INFO)
+    main()