FROG/load_train_args.py at main · RingBDStack/FROG · 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
141
142
143
144
145
146
147
# Due to the limitations of memory, multiple parameters need to be adjusted to ensure no OOM occurs
def load_train_args(args):
    dataset, task = args.dataset, args.task
    print(f'loading... | {dataset} | {task}')

    if dataset == 'rel-amazon':
        raise NotImplementedError("rel-amazon dataset is not supported.")

    elif dataset == 'rel-avito':
        if task == 'user-visits':
            args.batch_size = 1024
            args.dropout = 0.3
            args.lr = 0.001
            args.num_neighbors = 64
            args.channels = 32
        elif task == 'user-clicks':
            args.batch_size = 512
            args.dropout = 0.5
            args.lr = 0.0005
            args.num_layers = 1
            args.epochs = 20
            args.scheduler = True
        elif task == 'ad-ctr':
            args.batch_size = 128
            args.epochs = 20
            args.lr = 0.001
            args.num_neighbors = 64
            args.channels = 32
        elif task == 'user-ad-visit':
            args.batch_size = 256
            args.num_neighbors = 64
            args.num_layers = 8

    elif dataset == 'rel-event':
        args.batch_size = 128
        if task == 'user-repeat':
            args.dropout = 0.5
        elif task == 'user-ignore':
            args.lr = 0.0001
            args.dropout = 0.5
        elif task == 'user-attendance':
            args.epochs = 20
            args.MAE = 0.9999
            args.gamma = 0.

    elif dataset == 'rel-f1':
        args.batch_size = 64
        if task == 'driver-top3':
            args.num_layers = 1
            args.lr = 0.005
            args.dropout = 0.3 # 0.5
        elif task == 'driver-dnf':
            args.lr = 0.005
            args.epochs = 15
        elif task == 'driver-position':
            args.batch_size = 256
            args.num_layers = 1
            args.lr = 0.0005
            args.dropout = 0.5
            args.num_neighbors = 64

    elif dataset == 'rel-hm':
        if task == 'user-churn':
            args.channels = 32
            args.batch_size = 512
            args.lr = 0.001
            args.epochs = 20
            args.scheduler = True
            args.MAE = 0.95
        elif task == 'item-sales':
            args.batch_size = 256
            args.lr = 0.001
            args.epochs = 20
            args.scheduler = True
        elif task == 'user-item-purchase':
            args.num_neighbors = 64
            args.batch_size = 512
            args.channels = 64
            args.num_layers = 3
            args.lr = 0.0005
            args.dropout = 0.3

    elif dataset == 'rel-stack':
        args.num_neighbors = 64
        args.channels = 32
        if task == 'user-badge':
            args.batch_size = 1024
            args.lr = 0.005  #temp
            args.epochs = 20
            args.scheduler = True
        elif task == 'user-engagement':
            args.dropout = 0.3
            args.lr = 0.001
            args.batch_size = 1024
            args.epochs = 20
            args.scheduler = True
        elif task == 'post-votes':
            args.batch_size = 512
            args.lr = 0.0005
            args.epochs = 20
            args.scheduler = True
            args.warmup = 5
        elif task == 'user-post-comment':
            args.batch_size = 128
            args.channels = 64
            args.lr = 0.005
            args.num_layers = 6
            args.scheduler = True
        elif task == 'post-post-related':
            args.batch_size = 512
            args.channels = 64
            args.lr = 0.005
            args.dropout = 0.2

    elif dataset == 'rel-trial':
        if task == 'study-outcome':
            args.aggr = 'mean'
            args.batch_size = 256
            args.num_neighbors = 64
            args.channels = 64
            args.lr = 0.001
            args.dropout = 0.2
        elif task == 'site-success':
            args.aggr = 'mean'
            args.batch_size = 128
            args.channels = 64
            args.lr = 0.001
            args.dropout = 0.3
        elif task == 'study-adverse':
            args.batch_size = 128
            args.num_neighbors = 64
            args.channels = 64
            args.lr = 0.001
            args.dropout = 0.3
            args.epochs = 20
        elif task == 'condition-sponsor-run':
            args.batch_size = 128
            args.num_neighbors = 64
            args.channels = 128
            args.num_layers = 6
            args.lr = 0.0001
        elif task == 'site-sponsor-run':
            args.batch_size = 256
            args.num_neighbors = 64
            args.num_layers = 4
            args.lr = 0.005
    return args