-
Notifications
You must be signed in to change notification settings - Fork 3
/
train_meta.py
66 lines (54 loc) · 2.26 KB
/
train_meta.py
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
from deeprobust.graph.data import Pyg2Dpr
import argparse
import numpy as np
from meta_auto_ssl_homo import MetaSSL
from selfsl.gnn_encoder import GCN
from utils import *
import torch
import random
import time
st = time.time()
parser = argparse.ArgumentParser()
parser.add_argument('--gpu_id', type=int, default=0, help='gpu id')
parser.add_argument('--dataset', type=str, default='cora')
parser.add_argument('--epochs', type=int, default=1000)
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--lr_lambda', type=float, default=0.05)
parser.add_argument('--weight_decay', type=float, default=0.0000)
parser.add_argument('--hidden', type=int, default=512)
parser.add_argument('--dropout', type=float, default=0.0)
parser.add_argument('--normalize_features', type=bool, default=True)
parser.add_argument('--seed', type=int, default=15, help='Random seed.')
parser.add_argument('--fix_weight', type=int, default=0)
parser.add_argument('--num_layers', type=int, default=1)
parser.add_argument('--ssl', type=str)
parser.add_argument('--all_clusters', type=int, default=0)
args = parser.parse_args()
torch.cuda.set_device(args.gpu_id)
print('===========')
print(args)
# random seed setting
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
data = get_dataset(args.dataset, args.normalize_features)
nfeat = data.features.shape[1]
args.nclass = max(data.labels) + 1
set_of_ssl = ['Par', 'Clu', 'DGI', 'PairwiseDistance', 'PairwiseAttrSim']
if data.adj.shape[0] > 5000:
set_of_ssl = [ssl if ssl != 'DGI' else 'DGISample' for ssl in set_of_ssl]
else:
set_of_ssl = set_of_ssl
encoder = GCN(nfeat=nfeat, nhid=args.hidden, dropout=args.dropout, nlayers=args.num_layers, activation='prelu')
if args.dataset == 'arxiv':
encoder = GCN(nfeat=nfeat, nhid=args.hidden, dropout=args.dropout, nlayers=2, activation='prelu', with_bn=True, with_res=True)
auto_agent = MetaSSL(data, encoder, set_of_ssl, args, fix_weight=args.fix_weight)
args.epochs = 3000 if args.dataset in ['computers'] else 1000
x = auto_agent.pretrain(patience=1000)
print('Time: ', time.time()-st)
auto_agent.evaluate_pretrained(x)
# if args.dataset == 'arxiv':
# auto_agent.evaluate_pretrained_arxiv(x)
# else:
# auto_agent.evaluate_pretrained(x)