interaction_block.py 3.78 KB
Newer Older
1
2
import torch
import torch.nn as nn
3
4
5
from modules.initializers import GlorotOrthogonal
from modules.residual_layer import ResidualLayer

6
7
8
9
import dgl.function as fn


class InteractionBlock(nn.Module):
10
11
12
13
14
15
16
17
18
19
    def __init__(
        self,
        emb_size,
        num_radial,
        num_spherical,
        num_bilinear,
        num_before_skip,
        num_after_skip,
        activation=None,
    ):
20
21
22
23
24
        super(InteractionBlock, self).__init__()

        self.activation = activation
        # Transformations of Bessel and spherical basis representations
        self.dense_rbf = nn.Linear(num_radial, emb_size, bias=False)
25
26
27
        self.dense_sbf = nn.Linear(
            num_radial * num_spherical, num_bilinear, bias=False
        )
28
29
30
31
        # Dense transformations of input messages
        self.dense_ji = nn.Linear(emb_size, emb_size)
        self.dense_kj = nn.Linear(emb_size, emb_size)
        # Bilinear layer
32
33
34
        bilin_initializer = torch.empty(
            (emb_size, num_bilinear, emb_size)
        ).normal_(mean=0, std=2 / emb_size)
35
36
        self.W_bilin = nn.Parameter(bilin_initializer)
        # Residual layers before skip connection
37
38
39
40
41
42
        self.layers_before_skip = nn.ModuleList(
            [
                ResidualLayer(emb_size, activation=activation)
                for _ in range(num_before_skip)
            ]
        )
43
44
        self.final_before_skip = nn.Linear(emb_size, emb_size)
        # Residual layers after skip connection
45
46
47
48
49
50
        self.layers_after_skip = nn.ModuleList(
            [
                ResidualLayer(emb_size, activation=activation)
                for _ in range(num_after_skip)
            ]
        )
51
52

        self.reset_params()
53

54
55
56
57
58
59
60
61
62
    def reset_params(self):
        GlorotOrthogonal(self.dense_rbf.weight)
        GlorotOrthogonal(self.dense_sbf.weight)
        GlorotOrthogonal(self.dense_ji.weight)
        GlorotOrthogonal(self.dense_kj.weight)
        GlorotOrthogonal(self.final_before_skip.weight)

    def edge_transfer(self, edges):
        # Transform from Bessel basis to dence vector
63
        rbf = self.dense_rbf(edges.data["rbf"])
64
        # Initial transformation
65
66
        x_ji = self.dense_ji(edges.data["m"])
        x_kj = self.dense_kj(edges.data["m"])
67
68
69
70
71
        if self.activation is not None:
            x_ji = self.activation(x_ji)
            x_kj = self.activation(x_kj)

        # w: W * e_RBF \bigodot \sigma(W * m + b)
72
        return {"x_kj": x_kj * rbf, "x_ji": x_ji}
73
74

    def msg_func(self, edges):
75
        sbf = self.dense_sbf(edges.data["sbf"])
76
77
        # Apply bilinear layer to interactions and basis function activation
        # [None, 8] * [128, 8, 128] * [None, 128] -> [None, 128]
78
79
80
81
        x_kj = torch.einsum(
            "wj,wl,ijl->wi", sbf, edges.src["x_kj"], self.W_bilin
        )
        return {"x_kj": x_kj}
82
83
84

    def forward(self, g, l_g):
        g.apply_edges(self.edge_transfer)
85

86
87
88
89
90
        # nodes correspond to edges and edges correspond to nodes in the original graphs
        # node: d, rbf, o, rbf_env, x_kj, x_ji
        for k, v in g.edata.items():
            l_g.ndata[k] = v

91
        l_g.update_all(self.msg_func, fn.sum("x_kj", "m_update"))
92
93
94
95
96

        for k, v in l_g.ndata.items():
            g.edata[k] = v

        # Transformations before skip connection
97
        g.edata["m_update"] = g.edata["m_update"] + g.edata["x_ji"]
98
        for layer in self.layers_before_skip:
99
100
            g.edata["m_update"] = layer(g.edata["m_update"])
        g.edata["m_update"] = self.final_before_skip(g.edata["m_update"])
101
        if self.activation is not None:
102
            g.edata["m_update"] = self.activation(g.edata["m_update"])
103
104

        # Skip connection
105
        g.edata["m"] = g.edata["m"] + g.edata["m_update"]
106
107
108

        # Transformations after skip connection
        for layer in self.layers_after_skip:
109
            g.edata["m"] = layer(g.edata["m"])
110

111
        return g