YuWang0103/LGGM-Text2Graph

1import math2 3import torch4import torch.nn as nn5from torch.nn.modules.dropout import Dropout6from torch.nn.modules.linear import Linear7from torch.nn.modules.normalization import LayerNorm8from torch.nn import functional as F9from torch import Tensor10 11import utils12from diffusion import diffusion_utils13from models.layers import Xtoy, Etoy, masked_softmax14 15 16class XEyTransformerLayer(nn.Module):17    """ Transformer that updates node, edge and global features18        d_x: node features19        d_e: edge features20        dz : global features21        n_head: the number of heads in the multi_head_attention22        dim_feedforward: the dimension of the feedforward network model after self-attention23        dropout: dropout probablility. 0 to disable24        layer_norm_eps: eps value in layer normalizations.25    """26    def __init__(self, dx: int, de: int, dy: int, n_head: int, dim_ffX: int = 2048,27                 dim_ffE: int = 128, dim_ffy: int = 2048, dropout: float = 0.1,28                 layer_norm_eps: float = 1e-5, device=None, dtype=None) -> None:29        kw = {'device': device, 'dtype': dtype}30        super().__init__()31 32        self.self_attn = NodeEdgeBlock(dx, de, dy, n_head, **kw)33 34        self.linX1 = Linear(dx, dim_ffX, **kw)35        self.linX2 = Linear(dim_ffX, dx, **kw)36        self.normX1 = LayerNorm(dx, eps=layer_norm_eps, **kw)37        self.normX2 = LayerNorm(dx, eps=layer_norm_eps, **kw)38        self.dropoutX1 = Dropout(dropout)39        self.dropoutX2 = Dropout(dropout)40        self.dropoutX3 = Dropout(dropout)41 42        self.linE1 = Linear(de, dim_ffE, **kw)43        self.linE2 = Linear(dim_ffE, de, **kw)44        self.normE1 = LayerNorm(de, eps=layer_norm_eps, **kw)45        self.normE2 = LayerNorm(de, eps=layer_norm_eps, **kw)46        self.dropoutE1 = Dropout(dropout)47        self.dropoutE2 = Dropout(dropout)48        self.dropoutE3 = Dropout(dropout)49 50        self.lin_y1 = Linear(dy, dim_ffy, **kw)51        self.lin_y2 = Linear(dim_ffy, dy, **kw)52        self.norm_y1 = LayerNorm(dy, eps=layer_norm_eps, **kw)53        self.norm_y2 = LayerNorm(dy, eps=layer_norm_eps, **kw)54        self.dropout_y1 = Dropout(dropout)55        self.dropout_y2 = Dropout(dropout)56        self.dropout_y3 = Dropout(dropout)57 58        self.activation = F.relu59 60    def forward(self, X: Tensor, E: Tensor, y, node_mask: Tensor):61        """ Pass the input through the encoder layer.62            X: (bs, n, d)63            E: (bs, n, n, d)64            y: (bs, dy)65            node_mask: (bs, n) Mask for the src keys per batch (optional)66            Output: newX, newE, new_y with the same shape.67        """68 69        newX, newE, new_y = self.self_attn(X, E, y, node_mask=node_mask)70 71        newX_d = self.dropoutX1(newX)72        X = self.normX1(X + newX_d)73 74        newE_d = self.dropoutE1(newE)75        E = self.normE1(E + newE_d)76 77        new_y_d = self.dropout_y1(new_y)78        y = self.norm_y1(y + new_y_d)79 80        ff_outputX = self.linX2(self.dropoutX2(self.activation(self.linX1(X))))81        ff_outputX = self.dropoutX3(ff_outputX)82        X = self.normX2(X + ff_outputX)83 84        ff_outputE = self.linE2(self.dropoutE2(self.activation(self.linE1(E))))85        ff_outputE = self.dropoutE3(ff_outputE)86        E = self.normE2(E + ff_outputE)87 88        ff_output_y = self.lin_y2(self.dropout_y2(self.activation(self.lin_y1(y))))89        ff_output_y = self.dropout_y3(ff_output_y)90        y = self.norm_y2(y + ff_output_y)91 92        return X, E, y93 94 95class NodeEdgeBlock(nn.Module):96    """ Self attention layer that also updates the representations on the edges. """97    def __init__(self, dx, de, dy, n_head, **kwargs):98        super().__init__()99        assert dx % n_head == 0, f"dx: {dx} -- nhead: {n_head}"100        self.dx = dx101        self.de = de102        self.dy = dy103        self.df = int(dx / n_head)104        self.n_head = n_head105 106        # Attention107        self.q = Linear(dx, dx)108        self.k = Linear(dx, dx)109        self.v = Linear(dx, dx)110 111        # FiLM E to X112        self.e_add = Linear(de, dx)113        self.e_mul = Linear(de, dx)114 115        # FiLM y to E116        self.y_e_mul = Linear(dy, dx)           # Warning: here it's dx and not de117        self.y_e_add = Linear(dy, dx)118 119        # FiLM y to X120        self.y_x_mul = Linear(dy, dx)121        self.y_x_add = Linear(dy, dx)122 123        # Process y124        self.y_y = Linear(dy, dy)125        self.x_y = Xtoy(dx, dy)126        self.e_y = Etoy(de, dy)127 128        # Output layers129        self.x_out = Linear(dx, dx)130        self.e_out = Linear(dx, de)131        self.y_out = nn.Sequential(nn.Linear(dy, dy), nn.ReLU(), nn.Linear(dy, dy))132 133    def forward(self, X, E, y, node_mask):134        """135        :param X: bs, n, d        node features136        :param E: bs, n, n, d     edge features137        :param y: bs, dz           global features138        :param node_mask: bs, n139        :return: newX, newE, new_y with the same shape.140        """141        bs, n, _ = X.shape142        x_mask = node_mask.unsqueeze(-1)        # bs, n, 1143        e_mask1 = x_mask.unsqueeze(2)           # bs, n, 1, 1144        e_mask2 = x_mask.unsqueeze(1)           # bs, 1, n, 1145 146        # 1. Map X to keys and queries147        Q = self.q(X) * x_mask           # (bs, n, dx)148        K = self.k(X) * x_mask           # (bs, n, dx)149        diffusion_utils.assert_correctly_masked(Q, x_mask)150        # 2. Reshape to (bs, n, n_head, df) with dx = n_head * df151 152        Q = Q.reshape((Q.size(0), Q.size(1), self.n_head, self.df))153        K = K.reshape((K.size(0), K.size(1), self.n_head, self.df))154 155        Q = Q.unsqueeze(2)                              # (bs, 1, n, n_head, df)156        K = K.unsqueeze(1)                              # (bs, n, 1, n head, df)157 158        # Compute unnormalized attentions. Y is (bs, n, n, n_head, df)159        Y = Q * K160        Y = Y / math.sqrt(Y.size(-1))161        diffusion_utils.assert_correctly_masked(Y, (e_mask1 * e_mask2).unsqueeze(-1))162 163        E1 = self.e_mul(E) * e_mask1 * e_mask2                        # bs, n, n, dx164        E1 = E1.reshape((E.size(0), E.size(1), E.size(2), self.n_head, self.df))165 166        E2 = self.e_add(E) * e_mask1 * e_mask2                        # bs, n, n, dx167        E2 = E2.reshape((E.size(0), E.size(1), E.size(2), self.n_head, self.df))168 169        # Incorporate edge features to the self attention scores.170        Y = Y * (E1 + 1) + E2                  # (bs, n, n, n_head, df)171 172        # Incorporate y to E173        newE = Y.flatten(start_dim=3)                      # bs, n, n, dx174        ye1 = self.y_e_add(y).unsqueeze(1).unsqueeze(1)  # bs, 1, 1, de175        ye2 = self.y_e_mul(y).unsqueeze(1).unsqueeze(1)176        newE = ye1 + (ye2 + 1) * newE177 178        # Output E179        newE = self.e_out(newE) * e_mask1 * e_mask2      # bs, n, n, de180        diffusion_utils.assert_correctly_masked(newE, e_mask1 * e_mask2)181 182        # Compute attentions. attn is still (bs, n, n, n_head, df)183        softmax_mask = e_mask2.expand(-1, n, -1, self.n_head)    # bs, 1, n, 1184        attn = masked_softmax(Y, softmax_mask, dim=2)  # bs, n, n, n_head185 186        V = self.v(X) * x_mask                        # bs, n, dx187        V = V.reshape((V.size(0), V.size(1), self.n_head, self.df))188        V = V.unsqueeze(1)                                     # (bs, 1, n, n_head, df)189 190        # Compute weighted values191        weighted_V = attn * V192        weighted_V = weighted_V.sum(dim=2)193 194        # Send output to input dim195        weighted_V = weighted_V.flatten(start_dim=2)            # bs, n, dx196 197        # Incorporate y to X198        yx1 = self.y_x_add(y).unsqueeze(1)199        yx2 = self.y_x_mul(y).unsqueeze(1)200        newX = yx1 + (yx2 + 1) * weighted_V201 202        # Output X203        newX = self.x_out(newX) * x_mask204        diffusion_utils.assert_correctly_masked(newX, x_mask)205 206        # Process y based on X axnd E207        y = self.y_y(y)208        e_y = self.e_y(E)209        x_y = self.x_y(X)210        new_y = y + x_y + e_y211        new_y = self.y_out(new_y)               # bs, dy212 213        return newX, newE, new_y214 215 216class GraphTransformer(nn.Module):217    """218    n_layers : int -- number of layers219    dims : dict -- contains dimensions for each feature type220    """221    def __init__(self, n_layers: int, input_dims: dict, cond_dims: int, hidden_mlp_dims: dict, hidden_dims: dict,222                 output_dims: dict, act_fn_in: nn.ReLU(), act_fn_out: nn.ReLU()):223        super().__init__()224        self.n_layers = n_layers225        self.out_dim_X = output_dims['X']226        self.out_dim_E = output_dims['E']227        self.out_dim_y = output_dims['y']228 229        self.mlp_in_X = nn.Sequential(nn.Linear(input_dims['X'] + cond_dims, hidden_mlp_dims['X']), act_fn_in,230                                      nn.Linear(hidden_mlp_dims['X'], hidden_dims['dx']), act_fn_in)231 232        self.mlp_in_E = nn.Sequential(nn.Linear(input_dims['E'] + cond_dims, hidden_mlp_dims['E']), act_fn_in,233                                      nn.Linear(hidden_mlp_dims['E'], hidden_dims['de']), act_fn_in)234 235        self.mlp_in_y = nn.Sequential(nn.Linear(input_dims['y'], hidden_mlp_dims['y']), act_fn_in,236                                      nn.Linear(hidden_mlp_dims['y'], hidden_dims['dy']), act_fn_in)237 238        self.tf_layers = nn.ModuleList([XEyTransformerLayer(dx=hidden_dims['dx'],239                                                            de=hidden_dims['de'],240                                                            dy=hidden_dims['dy'],241                                                            n_head=hidden_dims['n_head'],242                                                            dim_ffX=hidden_dims['dim_ffX'],243                                                            dim_ffE=hidden_dims['dim_ffE'])244                                        for i in range(n_layers)])245 246        self.mlp_out_X = nn.Sequential(nn.Linear(hidden_dims['dx'], hidden_mlp_dims['X']), act_fn_out,247                                       nn.Linear(hidden_mlp_dims['X'], output_dims['X']))248 249        self.mlp_out_E = nn.Sequential(nn.Linear(hidden_dims['de'], hidden_mlp_dims['E']), act_fn_out,250                                       nn.Linear(hidden_mlp_dims['E'], output_dims['E']))251 252        self.mlp_out_y = nn.Sequential(nn.Linear(hidden_dims['dy'], hidden_mlp_dims['y']), act_fn_out,253                                       nn.Linear(hidden_mlp_dims['y'], output_dims['y']))254 255    def forward(self, X, E, y, node_mask):256        bs, n = X.shape[0], X.shape[1]257 258        diag_mask = torch.eye(n)259        diag_mask = ~diag_mask.type_as(E).bool()260        diag_mask = diag_mask.unsqueeze(0).unsqueeze(-1).expand(bs, -1, -1, -1)261 262        X_to_out = X[..., :self.out_dim_X]263        E_to_out = E[..., :self.out_dim_E]264        y_to_out = y[..., :self.out_dim_y]265 266        new_E = self.mlp_in_E(E)267        new_E = (new_E + new_E.transpose(1, 2)) / 2268        269        after_in = utils.PlaceHolder(X=self.mlp_in_X(X), E=new_E, y=self.mlp_in_y(y)).mask(node_mask)270        X, E, y = after_in.X, after_in.E, after_in.y271 272        for layer in self.tf_layers:273            X, E, y = layer(X, E, y, node_mask)274 275        X = self.mlp_out_X(X)276        E = self.mlp_out_E(E)277        y = self.mlp_out_y(y)278 279        X = (X + X_to_out)280        E = (E + E_to_out) * diag_mask281        y = y + y_to_out282 283        E = 1/2 * (E + torch.transpose(E, 1, 2))284 285        return utils.PlaceHolder(X=X, E=E, y=y).mask(node_mask)286