src/models/CatGCN/layers.py

import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from torch.nn.parameter import Parameter

import numpy as np
import networkx as nx
from torch_geometric.nn import GCNConv, GATConv, SGConv, APPNP

#from gnn_layers import BatchAGC, BatchFiGNN, BatchGAT
from CatGCN.gnn_layers import BatchAGC, BatchFiGNN, BatchGAT
#from pna_layer import PNAConv
from CatGCN.pna_layer import PNAConv
#from gcnii_layer import GCNIIConv
from CatGCN.gcnii_layer import GCNIIConv

class StackedGNN(nn.Module):
    """
    Multi-layer GNN model.
    """
    def __init__(self, args, field_count, field_size, output_channels):
        """
        :param args: Arguments object.
        :param field_count: Number of fields.
        :param field_size: Number of sampled fields for each user.
        :param output_channels: Number of target classes.
        """
        super(StackedGNN, self).__init__()
        self.args = args

        if self.args.grn_units != 'none':
            self.grn_units = [args.field_dim] + [int(x) for x in args.grn_units.strip().split(",")] + [output_channels]
        else:
            self.grn_units = [args.field_dim] + [output_channels]
        if self.args.nfm_units != 'none':
            self.nfm_units = [args.field_dim] + [int(x) for x in args.nfm_units.strip().split(",")] + [output_channels]
        else:
            self.nfm_units = [args.field_dim] + [output_channels]

        self.input_channels = args.field_dim
        self.output_channels = output_channels

        # For Baseline
        if self.args.gnn_units != 'none':
            self.gnn_units = [self.input_channels] + [int(x) for x in args.gnn_units.strip().split(",")] + [self.output_channels]
        else:
            self.gnn_units = [self.input_channels] + [self.output_channels]

        self.field_count = field_count
        self.field_size = field_size


        self.field_embedding = nn.Embedding(field_count, args.field_dim)
        self.field_embedding.weight.requires_grad = True
                
        self._setup_layers()

    def _setup_layers(self):
        """
        Creating the layers based on the args.
        """
        # Categorical feature interaction modeling
        ''' Global interaction modeling '''
        if self.args.graph_refining == 'agc':
            self.grn = BatchAGC(self.args.field_dim, self.args.field_dim)
            self.num_grn_layer = len(self.grn_units) - 1
            self.grn_layer_stack = nn.ModuleList()
            for i in range(self.num_grn_layer):
                self.grn_layer_stack.append(
                        nn.Linear(self.grn_units[i], self.grn_units[i + 1], bias=True))
        elif self.args.graph_refining == 'gat':
            n_heads = [int(x) for x in self.args.multi_heads.strip().split(",")]
            attn_dropout = 0. 
            # attn_dropout = self.args.dropout
            self.gat_units = [int(x) for x in self.args.gat_units.strip().split(",")]
            self.num_gat_layer = len(self.gat_units) - 1
            self.gat_layer_stack = nn.ModuleList()
            for i in range(self.num_gat_layer):
                f_in = self.gat_units[i] * n_heads[i - 1] if i else self.gat_units[i] 
                self.gat_layer_stack.append(
                        BatchGAT(
                            n_heads[i], f_in=f_in,
                            f_out=self.gat_units[i + 1], attn_dropout=attn_dropout))
            self.num_grn_layer = len(self.grn_units) - 1
            self.grn_layer_stack = nn.ModuleList()
            for i in range(self.num_grn_layer):
                self.grn_layer_stack.append(
                        nn.Linear(self.grn_units[i], self.grn_units[i + 1], bias=True))
        elif self.args.graph_refining == 'cosimi':
            self.num_grn_layer = len(self.grn_units) - 1
            self.grn_layer_stack = nn.ModuleList()
            for i in range(self.num_grn_layer):
                self.grn_layer_stack.append(
                        nn.Linear(self.grn_units[i], self.grn_units[i + 1], bias=True))

        ''' Local interaction modeling '''
        if self.args.bi_interaction == 'nfm':
            self.num_nfm_layer = len(self.nfm_units) - 1
            self.nfm_layer_stack = nn.ModuleList()
            for i in range(self.num_nfm_layer):
                self.nfm_layer_stack.append(
                        nn.Linear(self.nfm_units[i], self.nfm_units[i + 1], bias=True))

        # GNN Layer
        if self.args.graph_layer == 'gcn':
            self.gnn_layers = nn.ModuleList()
            for i, _ in enumerate(self.gnn_units[:-1]):
                self.gnn_layers.append(GCNConv(self.gnn_units[i], self.gnn_units[i+1]))
        elif self.args.graph_layer == 'gat_1':
            self.gnn_layers = GATConv(self.input_channels, self.output_channels, heads=1, concat=True, negative_slope=0.2, dropout=self.args.dropout, bias=True)
        elif self.args.graph_layer == 'gat_2':
            n_heads = 8
            self.gnn_layers_1 = GATConv(self.input_channels, self.gnn_units[1], heads=n_heads, concat=True, negative_slope=0.2, dropout=0, bias=True)
            self.gnn_layers_2 = GATConv(self.gnn_units[1]*n_heads, self.output_channels, heads=1, concat=True, negative_slope=0.2, dropout=0, bias=True)
        elif self.args.graph_layer == 'sgc':
            self.gnn_layers = SGConv(self.input_channels, self.output_channels, K=self.args.gnn_hops, cached=False)
        elif self.args.graph_layer == 'appnp':
            self.num_mlp_layer = len(self.gnn_units) - 1
            self.mlp_layer_stack = nn.ModuleList()
            for i in range(self.num_mlp_layer):
                self.mlp_layer_stack.append(
                        nn.Linear(self.gnn_units[i], self.gnn_units[i + 1], bias=True))
            self.gnn_layers = APPNP(K=10, alpha=0.1, bias=True)
        elif self.args.graph_layer == 'cat-appnp':
            self.gnn_layers = APPNP(K=10, alpha=0.1, bias=True)
        elif self.args.graph_layer == 'gcnii_F':
            self.num_gnn_layer = self.args.gnn_hops
            self.lin_layer_1 = nn.Linear(self.input_channels, self.gnn_units[1], bias=True)
            self.gnn_layers = nn.ModuleList()
            for layer in range(self.num_gnn_layer):
                self.gnn_layers.append(GCNIIConv(self.gnn_units[1], alpha=self.args.alpha, theta=self.args.theta, layer=layer+1, shared_weights=False))
            self.lin_layer_2 = nn.Linear(self.gnn_units[1], self.output_channels, bias=True)
        elif self.args.graph_layer == 'gcnii_T':
            self.num_gnn_layer = self.args.gnn_hops
            self.lin_layer_1 = nn.Linear(self.input_channels, self.gnn_units[1], bias=True)
            self.gnn_layers = nn.ModuleList()
            for layer in range(self.num_gnn_layer):
                self.gnn_layers.append(GCNIIConv(self.gnn_units[1], alpha=self.args.alpha, theta=self.args.theta, layer=layer+1, shared_weights=True))
            self.lin_layer_2 = nn.Linear(self.gnn_units[1], self.output_channels, bias=True)
        elif self.args.graph_layer == 'cross_1':
            self.mlp_layers_1 = nn.Linear(self.input_channels, self.output_channels, bias=False)
            self.mlp_layers_2 = nn.Linear(self.input_channels, self.output_channels, bias=False)
            self.gnn_layers = PNAConv(K=1, cached=False)
        elif self.args.graph_layer == 'cross_2':
            self.mlp_layers_11 = nn.Linear(self.input_channels, self.gnn_units[1], bias=False)
            self.mlp_layers_12 = nn.Linear(self.input_channels, self.gnn_units[1], bias=False)
            self.gnn_layers_1 = PNAConv(K=1, cached=False)
            self.mlp_layers_21 = nn.Linear(self.gnn_units[1], self.output_channels, bias=False)
            self.mlp_layers_22 = nn.Linear(self.gnn_units[1], self.output_channels, bias=False)
            self.gnn_layers_2 = PNAConv(K=1, cached=False)
        elif self.args.graph_layer == 'fignn':
            self.fi_layers = BatchFiGNN(self.input_channels, self.gnn_units[1], self.output_channels)
            self.gnn_layers = PNAConv(K=self.args.gnn_hops, cached=False)
        elif self.args.graph_layer == 'pna':
            self.gnn_layers = PNAConv(K=self.args.gnn_hops, cached=False)

    def forward(self, edges, field_index, field_adjs):
        """
        Making a forward pass.
        :param edges: Edge list LongTensor.
        :parm field_index: User-field index matrix.
        :parm field_adjs: Normalized adjacency matrix with probe coefficient.
        :return predictions: Prediction matrix output FLoatTensor.
        """
        raw_field_feature = self.field_embedding(field_index)

        # Categorical feature interaction modeling
        ''' Global interaction modeling '''
        field_feature = raw_field_feature
        
        if self.args.graph_refining == 'agc':
            field_feature = self.grn(field_feature, field_adjs.float())
            field_feature = F.relu(field_feature)
            field_feature = F.dropout(field_feature, self.args.dropout, training=self.training)
            
            if self.args.aggr_pooling == 'mean':
                user_feature = torch.mean(field_feature, dim=-2)

            for i, grn_layer in enumerate(self.grn_layer_stack):
                user_feature = grn_layer(user_feature)
                if i + 1 < self.num_grn_layer:
                    user_feature = F.relu(user_feature)
                    user_feature = F.dropout(user_feature, self.args.dropout, training=self.training)
            user_gnn_feature = user_feature
        
        elif self.args.graph_refining == 'gat':
            bs, n = field_adjs.size()[:2]
            for i, gat_layer in enumerate(self.gat_layer_stack):
                field_feature = gat_layer(field_feature, field_adjs.byte()) 
                if i + 1 == self.num_gat_layer:
                    field_feature = field_feature.mean(dim=1) 
                else:
                    field_feature = F.elu(field_feature.transpose(1, 2).contiguous().view(bs, n, -1)) 
                    field_feature = F.dropout(field_feature, self.args.dropout, training=self.training)

            if self.args.aggr_pooling == 'mean':
                user_feature = torch.mean(field_feature, dim=-2)
            for i, grn_layer in enumerate(self.grn_layer_stack):
                user_feature = grn_layer(user_feature)
                if i + 1 < self.num_grn_layer:
                    user_feature = F.relu(user_feature)
                    user_feature = F.dropout(user_feature, self.args.dropout, training=self.training)
            user_gnn_feature = user_feature
        
        elif self.args.graph_refining == 'cosimi':
            similarity_mat = torch.bmm(field_feature, field_feature.permute(0, 2, 1)) 
            feature_norm = torch.sqrt(torch.sum(torch.mul(field_feature, field_feature), dim=-1)).unsqueeze(2) 
            cosine_distance = torch.div(similarity_mat, torch.mul(feature_norm, feature_norm.permute(0, 2, 1))) 
            field_feature = torch.bmm(cosine_distance, field_feature) 

            if self.args.aggr_pooling == 'mean':
                user_feature = torch.mean(field_feature, dim=-2)

            for i, grn_layer in enumerate(self.grn_layer_stack):
                user_feature = grn_layer(user_feature)
                if i + 1 < self.num_grn_layer:
                    user_feature = F.relu(user_feature)
                    user_feature = F.dropout(user_feature, self.args.dropout, training=self.training)
            user_gnn_feature = user_feature

        ''' Local interaction modeling '''
        field_feature = raw_field_feature
        if self.args.bi_interaction == 'nfm': 
            # sum-square-part
            summed_field_feature = torch.sum(field_feature, 1) 
            square_summed_field_feature = summed_field_feature ** 2 
            # squre-sum-part
            squared_field_feature = field_feature ** 2 
            sum_squared_field_feature = torch.sum(squared_field_feature, 1) 
            # second order
            user_feature = 0.5 * (square_summed_field_feature - sum_squared_field_feature)
            # deep part
            for i, nfm_layer in enumerate(self.nfm_layer_stack):
                user_feature = nfm_layer(user_feature)
                if i + 1 < self.num_nfm_layer:
                    user_feature = F.relu(user_feature)
                    user_feature = F.dropout(user_feature, self.args.dropout, training=self.training)
            user_nfm_feature = user_feature

        # Aggregation
        if self.args.aggr_style == 'sum':
            user_feature = self.args.balance_ratio*user_gnn_feature + \
                (1-self.args.balance_ratio)*user_nfm_feature 

        if self.args.graph_refining == 'none' and self.args.bi_interaction == 'none':
            user_feature = torch.mean(raw_field_feature, dim=-2)

        # GNN Layer
        if self.args.graph_layer == 'gcn':
            for i, _ in enumerate(self.gnn_units[:-2]): 
                user_feature = F.relu(self.gnn_layers[i](user_feature, edges))
                if i > 1:
                    user_feature = F.dropout(user_feature, p=self.args.dropout, training=self.training)
            user_feature = self.gnn_layers[i+1](user_feature, edges)
            predictions = F.log_softmax(user_feature, dim=1)
        elif self.args.graph_layer == 'gat_1': 
            user_feature = self.gnn_layers(user_feature, edges)
            predictions = F.log_softmax(user_feature, dim=1)
        elif self.args.graph_layer == 'gat_2': 
            user_feature = F.elu(self.gnn_layers_1(user_feature, edges))
            user_feature = F.dropout(user_feature, p=self.args.dropout, training=self.training)
            user_feature = self.gnn_layers_2(user_feature, edges)
            predictions = F.log_softmax(user_feature, dim=1)
        elif self.args.graph_layer == 'sgc':
            user_feature = self.gnn_layers(user_feature, edges)
            predictions = F.log_softmax(user_feature, dim=1)
        elif self.args.graph_layer == 'appnp':
            for i, mlp_layer in enumerate(self.mlp_layer_stack):
                user_feature = mlp_layer(user_feature)
                if i + 1 < self.num_mlp_layer:
                    user_feature = F.relu(user_feature)
                    user_feature = F.dropout(user_feature, self.args.dropout, training=self.training)
            user_feature = self.gnn_layers(user_feature, edges)
            predictions = F.log_softmax(user_feature, dim=1)
        elif self.args.graph_layer == 'cat-appnp':
            user_feature = self.gnn_layers(user_feature, edges)
            predictions = F.log_softmax(user_feature, dim=1) 
        elif self.args.graph_layer == 'gcnii_F' or self.args.graph_layer == 'gcnii_T':
            user_feature = self.lin_layer_1(user_feature)
            user_feature = F.relu(user_feature)
            user_feature = user_feature_0 = F.dropout(user_feature, self.args.dropout, training=self.training)
            for i, gnn_layer in enumerate(self.gnn_layers):
                    user_feature = gnn_layer(user_feature, user_feature_0, edges)
                    if i + 1 < self.num_gnn_layer:
                        user_feature = F.relu(user_feature)
                        user_feature = F.dropout(user_feature, self.args.dropout, training=self.training)
            user_feature = self.lin_layer_2(user_feature)
            predictions = F.log_softmax(user_feature, dim=1)
        elif self.args.graph_layer == 'cross_1':
            alpha = 1
            user_feature = F.dropout(user_feature, self.args.dropout, training=self.training)
            x_1 = self.mlp_layers_1(user_feature)
            x_2 = self.mlp_layers_1(user_feature)
            x_sec_ord = torch.mul(x_1, x_2) * alpha 
            user_feature = x_1 + x_2 + x_sec_ord
            user_feature = self.gnn_layers(user_feature, edges)
            predictions = F.log_softmax(user_feature, dim=1)
        elif self.args.graph_layer == 'cross_2':
            alpha = 1
            user_feature = F.dropout(user_feature, self.args.dropout, training=self.training)
            x_11 = self.mlp_layers_11(user_feature)
            x_12 = self.mlp_layers_12(user_feature)
            x_sec_ord_1 = torch.mul(x_11, x_12) * alpha 
            user_feature = x_11 + x_12 + x_sec_ord_1
            user_feature = self.gnn_layers_1(user_feature, edges)
            user_feature = F.dropout(user_feature, self.args.dropout, training=self.training)
            x_21 = self.mlp_layers_21(user_feature)
            x_22 = self.mlp_layers_22(user_feature)
            x_sec_ord_2 = torch.mul(x_21, x_22) * alpha 
            user_feature = x_21 + x_22 + x_sec_ord_2
            user_feature = self.gnn_layers_2(user_feature, edges)
            predictions = F.log_softmax(user_feature, dim=1)
        elif self.args.graph_layer == 'fignn':
            user_feature = self.fi_layers(raw_field_feature, field_adjs.float(), self.args.num_steps)
            user_feature = self.gnn_layers(user_feature, edges)
            predictions = F.log_softmax(user_feature, dim=1)
        elif self.args.graph_layer == 'pna':
            user_feature = self.gnn_layers(user_feature, edges)
            #print('user_feature pna in forward pass:', user_feature)
            predictions = F.log_softmax(user_feature, dim=1)
        elif self.args.graph_layer == 'none':
            predictions = F.log_softmax(user_feature, dim=1)
        return predictions