ProG/pre_train.py at main · sheldonresearch/ProG · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
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
131
132
133
134
135
136
137
138
139
140
import argparse

from prompt_graph.data import load4graph
from prompt_graph.defines import GRAPH_TASKS, NODE_TASKS
from prompt_graph.pretrain import (
    DGI,
    Edgepred_GPPT,
    Edgepred_Gprompt,
    GraphCL,
    GraphMAE,
    NodePrePrompt,
    SimGRACE,
)
from prompt_graph.pretrain.MultiGPrompt import GraphPrePrompt
from prompt_graph.utils import get_args, resolve_device, seed_everything


def get_pretrain_task_by_dataset_name(dataset_name: str) -> str:
    if dataset_name in GRAPH_TASKS:
        return "GraphTask"
    elif dataset_name in NODE_TASKS:
        return "NodeTask"
    else:
        raise ValueError(f"Does not support this kind of dataset {dataset_name}.")


def get_pretrain_task_delegate(args: argparse.Namespace):
    seed_everything(args.seed)
    args.device = resolve_device(args.device)
    if args.pretrain_task == "SimGRACE":
        pt = SimGRACE(
            dataset_name=args.dataset_name,
            gnn_type=args.gnn_type,
            hid_dim=args.hid_dim,
            gln=args.num_layer,
            num_epoch=args.epochs,
            device=args.device,
            num_workers=args.num_workers,
        )
    elif args.pretrain_task == "GraphCL":
        pt = GraphCL(
            dataset_name=args.dataset_name,
            gnn_type=args.gnn_type,
            hid_dim=args.hid_dim,
            gln=args.num_layer,
            num_epoch=args.epochs,
            device=args.device,
            num_workers=args.num_workers,
        )
    elif args.pretrain_task == "Edgepred_GPPT":
        pt = Edgepred_GPPT(
            dataset_name=args.dataset_name,
            gnn_type=args.gnn_type,
            hid_dim=args.hid_dim,
            gln=args.num_layer,
            num_epoch=args.epochs,
            device=args.device,
            num_workers=args.num_workers,
        )
    elif args.pretrain_task == "Edgepred_Gprompt":
        pt = Edgepred_Gprompt(
            dataset_name=args.dataset_name,
            gnn_type=args.gnn_type,
            hid_dim=args.hid_dim,
            gln=args.num_layer,
            num_epoch=args.epochs,
            device=args.device,
            num_workers=args.num_workers,
        )
    elif args.pretrain_task == "DGI":
        pt = DGI(
            dataset_name=args.dataset_name,
            gnn_type=args.gnn_type,
            hid_dim=args.hid_dim,
            gln=args.num_layer,
            num_epoch=args.epochs,
            device=args.device,
            num_workers=args.num_workers,
        )
    elif args.pretrain_task in ("NodeMultiGprompt", "MultiGprompt", "GraphMultiGprompt"):
        if args.pretrain_task == "NodeMultiGprompt" or args.dataset_name in NODE_TASKS:
            nonlinearity = "prelu"
            pt = NodePrePrompt(
                args.dataset_name,
                args.hid_dim,
                nonlinearity,
                0.9,
                0.9,
                0.1,
                0.001,
                1,
                0.3,
                device=args.device,
            )
        elif args.pretrain_task == "GraphMultiGprompt" or args.dataset_name in GRAPH_TASKS:
            nonlinearity = "prelu"
            input_dim, out_dim, graph_list = load4graph(args.dataset_name, pretrained=True)
            pt = GraphPrePrompt(
                graph_list,
                input_dim,
                out_dim,
                args.dataset_name,
                args.hid_dim,
                nonlinearity,
                0.9,
                0.9,
                0.1,
                args.num_layer,
                0.3,
                args.device,
            )
        else:
            raise ValueError(
                f"Unsupported args.pretrain_task type for MultiGprompt {args.pretrain_task}"
            )
    elif args.pretrain_task == "GraphMAE":
        pt = GraphMAE(
            dataset_name=args.dataset_name,
            gnn_type=args.gnn_type,
            hid_dim=args.hid_dim,
            gln=args.num_layer,
            num_epoch=args.epochs,
            device=args.device,
            mask_rate=0.75,
            drop_edge_rate=0.0,
            replace_rate=0.1,
            loss_fn="sce",
            alpha_l=2,
            num_workers=args.num_workers,
        )
    else:
        raise ValueError(f"Unexpected args.pretrain_task type: {args.pretrain_task}.")
    return pt


if __name__ == "__main__":
    args = get_args()

    pt = get_pretrain_task_delegate(args=args)
    pt.pretrain()