DiffPool/tst.py at master · thomast8/DiffPool · 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
from core.no_cheb.helper import *
from core.no_cheb.layers import *
from core.no_cheb.training import make_batch

from tensorflow.keras.layers import Dense


def main():
    batch_size = 2

    dataset = read_graphfile("datasets", "ENZYMES", max_nodes=None)

    train, val, test, num_classes = make_train_test_val(dataset)
    avg_num_nodes = calculate_avg_nodes(dataset)

    train, val, test = preprocess_dataset(train, val, test)

    batch = make_batch(train, batch_size)

    for x, a, y, ind in batch:

        test = GCN(features=4, dropout=0.5)((a, x, ind))
        test = DiffPool(max_clusters=1)(test)
        # print("A:")
        # tf.print(test[0], summarize=-1)
        # print("data:")
        # tf.print(test[1], summarize=-1)
        # print("indicator:")
        # tf.print(test[2], summarize=-1)

        test = ReshapeForDense()(test)
        print("data reshaped:")
        tf.print(test, summarize=-1)

        test = Dense(num_classes, activation="softmax")(test)
        print("Classification:")
        tf.print(test, summarize=-1)

        return


main()