Exemple #1
0
class RGCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout, num_nodes_dict, x_types, edge_types):
        super(RGCN, self).__init__()

        node_types = list(num_nodes_dict.keys())

        self.embs = ParameterDict({
            key: Parameter(torch.Tensor(num_nodes_dict[key], in_channels))
            for key in set(node_types).difference(set(x_types))
        })

        self.convs = ModuleList()
        self.convs.append(
            RGCNConv(in_channels, hidden_channels, node_types, edge_types))
        for _ in range(num_layers - 2):
            self.convs.append(
                RGCNConv(hidden_channels, hidden_channels, node_types,
                         edge_types))
        self.convs.append(
            RGCNConv(hidden_channels, out_channels, node_types, edge_types))

        self.dropout = dropout

        self.reset_parameters()

    def reset_parameters(self):
        for emb in self.embs.values():
            torch.nn.init.xavier_uniform_(emb)
        for conv in self.convs:
            conv.reset_parameters()

    def forward(self, x_dict, adj_t_dict):
        x_dict = copy.copy(x_dict)
        for key, emb in self.embs.items():
            x_dict[key] = emb

        for conv in self.convs[:-1]:
            x_dict = conv(x_dict, adj_t_dict)
            for key, x in x_dict.items():
                x_dict[key] = F.relu(x)
                x_dict[key] = F.dropout(x,
                                        p=self.dropout,
                                        training=self.training)
        return self.convs[-1](x_dict, adj_t_dict)
Exemple #2
0
class RGCN(torch.nn.Module):
    def __init__(
        self,
        in_channels,
        hidden_channels,
        out_channels,
        num_layers,
        dropout,
        num_nodes_dict,
        x_types,
        num_edge_types,
    ):
        super(RGCN, self).__init__()

        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.out_channels = out_channels
        self.num_layers = num_layers
        self.dropout = dropout

        node_types = list(num_nodes_dict.keys())
        num_node_types = len(node_types)

        self.num_node_types = num_node_types
        self.num_edge_types = num_edge_types

        # Create embeddings for all node types that do not come with features.
        self.emb_dict = ParameterDict({
            f"{key}": Parameter(torch.Tensor(num_nodes_dict[key], in_channels))
            for key in set(node_types).difference(set(x_types))
        })

        I, H, O = in_channels, hidden_channels, out_channels  # noqa

        # Create `num_layers` many message passing layers.
        self.convs = ModuleList()
        self.convs.append(RGCNConv(I, H, num_node_types, num_edge_types))
        for _ in range(num_layers - 2):
            self.convs.append(RGCNConv(H, H, num_node_types, num_edge_types))
        self.convs.append(RGCNConv(H, O, self.num_node_types, num_edge_types))

        self.reset_parameters()

    def reset_parameters(self):
        for emb in self.emb_dict.values():
            torch.nn.init.xavier_uniform_(emb)
        for conv in self.convs:
            conv.reset_parameters()

    def group_input(self, x_dict, node_type, local_node_idx):
        # Create global node feature matrix.
        h = torch.zeros((node_type.size(0), self.in_channels),
                        device=node_type.device)

        for key, x in x_dict.items():
            mask = node_type == key
            h[mask] = x[local_node_idx[mask]]

        for key, emb in self.emb_dict.items():
            mask = node_type == int(key)
            h[mask] = emb[local_node_idx[mask]]

        return h

    def forward(self, x_dict, edge_index, edge_type, node_type,
                local_node_idx):

        x = self.group_input(x_dict, node_type, local_node_idx)

        for i, conv in enumerate(self.convs):
            x = conv(x, edge_index, edge_type, node_type)
            if i != self.num_layers - 1:
                x = F.relu(x)
                x = F.dropout(x, p=0.5, training=self.training)

        return x.log_softmax(dim=-1)

    def inference(self, x_dict, edge_index_dict, key2int):
        # We can perform full-batch inference on GPU.

        device = list(x_dict.values())[0].device

        x_dict = copy(x_dict)
        for key, emb in self.emb_dict.items():
            x_dict[int(key)] = emb

        adj_t_dict = {}
        for key, (row, col) in edge_index_dict.items():
            adj_t_dict[key] = SparseTensor(row=col, col=row).to(device)

        for i, conv in enumerate(self.convs):
            out_dict = {}

            for j, x in x_dict.items():
                out_dict[j] = conv.root_lins[j](x)

            for keys, adj_t in adj_t_dict.items():
                src_key, target_key = keys[0], keys[-1]
                out = out_dict[key2int[target_key]]
                tmp = adj_t.matmul(x_dict[key2int[src_key]], reduce="mean")
                out.add_(conv.rel_lins[key2int[keys]](tmp))

            if i != self.num_layers - 1:
                for j in range(self.num_node_types):
                    F.relu_(out_dict[j])

            x_dict = out_dict

        return x_dict
Exemple #3
0
class RGNN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout, num_nodes_dict, x_types, num_edge_types, args):
        super(RGNN, self).__init__()

        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.out_channels = out_channels
        self.num_layers = num_layers
        self.dropout = dropout
        self.args = args

        node_types = list(num_nodes_dict.keys())
        num_node_types = len(node_types)

        self.num_node_types = num_node_types
        self.num_edge_types = num_edge_types

        # Create embeddings for all node types that do not come with features.
        self.emb_dict = ParameterDict({
            f'{key}': Parameter(torch.Tensor(num_nodes_dict[key], in_channels))
            for key in set(node_types).difference(set(x_types))
        })

        I, H, O = in_channels, hidden_channels, out_channels  # noqa

        if self.args.conv_name == 'rgcn':
            self.convs = ModuleList()
            self.convs.append(
                RGCNConv(I, H, num_node_types, num_edge_types, self.args))
            for _ in range(num_layers - 2):
                self.convs.append(
                    RGCNConv(H, H, num_node_types, num_edge_types, self.args))
            self.convs.append(
                RGCNConv(H, O, self.num_node_types, num_edge_types, self.args))
        else:
            self.convs = ModuleList()
            self.convs.append(
                RGSNConv(I, H, num_node_types, num_edge_types, self.args))
            for _ in range(num_layers - 2):
                self.convs.append(
                    RGSNConv(H, H, num_node_types, num_edge_types, self.args))
            self.convs.append(
                RGSNConv(H, O, self.num_node_types, num_edge_types, self.args))

        if self.args.Norm4:
            self.norm = torch.nn.LayerNorm(I)

        self.reset_parameters()

    def reset_parameters(self):
        root = self.args.feat_dir
        Feat_list = [
            os.path.join(root, i) for i in [
                './author_FEAT.npy', './field_of_study_FEAT.npy',
                './institution_FEAT.npy'
            ]
        ]

        for emb, Feat_path in zip(self.emb_dict.values(), Feat_list):
            if self.args.FDFT:
                emb.data = torch.Tensor(np.load(Feat_path)).to(
                    self.args.device)
                emb.requires_grad = True
            else:
                torch.nn.init.xavier_uniform_(emb)

        for conv in self.convs:
            conv.reset_parameters()

        if self.args.Norm4:
            self.norm.reset_parameters()

    def group_input(self,
                    x_dict,
                    node_type,
                    local_node_idx,
                    n_id=None,
                    perturb=None):
        # Create global node feature matrix.
        if n_id is not None:
            node_type = node_type[n_id]
            local_node_idx = local_node_idx[n_id]

        h = torch.zeros((node_type.size(0), self.in_channels),
                        device=node_type.device)

        for key, x in x_dict.items():
            mask = node_type == int(key)
            h[mask] = x[local_node_idx[mask]] if perturb is None else x[
                local_node_idx[mask]] + perturb

        for key, emb in self.emb_dict.items():
            mask = node_type == int(key)
            h[mask] = emb[local_node_idx[mask]]

        return h

    def forward(self,
                n_id,
                x_dict,
                adjs,
                edge_type,
                node_type,
                local_node_idx,
                perturb=None):

        x = self.group_input(x_dict, node_type, local_node_idx, n_id, perturb)

        if self.args.FDFT:
            x = F.dropout(x, p=0.5, training=self.training)

        if self.args.Norm4:
            x = self.norm(x)

        node_type = node_type[n_id]

        for i, (edge_index, e_id, size) in enumerate(adjs):
            x_target = x[:size[1]]  # Target node embeddings.
            src_node_type = node_type
            node_type = node_type[:size[1]]  # Target node types.

            conv = self.convs[i]
            x = conv((x, x_target), edge_index, edge_type[e_id], node_type,
                     src_node_type)
            if i != self.num_layers - 1:
                x = F.relu(x)
                x = F.dropout(x, p=0.5, training=self.training)

        return x.log_softmax(dim=-1)
Exemple #4
0
class RGCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout, num_nodes_dict, x_types, num_edge_types):
        super(RGCN, self).__init__()

        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.out_channels = out_channels
        self.num_layers = num_layers
        self.dropout = dropout

        node_types = list(num_nodes_dict.keys())
        num_node_types = len(node_types)

        self.num_node_types = num_node_types
        self.num_edge_types = num_edge_types

        # Create embeddings for all node types that do not come with features.
        self.emb_dict = ParameterDict({
            f'{key}': Parameter(torch.Tensor(num_nodes_dict[key], in_channels))
            for key in set(node_types).difference(set(x_types))
        })

        I, H, O = in_channels, hidden_channels, out_channels  # noqa

        # Create `num_layers` many message passing layers.
        self.convs = ModuleList()
        self.convs.append(RGCNConv(I, H, num_node_types, num_edge_types))
        for _ in range(num_layers - 2):
            self.convs.append(RGCNConv(H, H, num_node_types, num_edge_types))
        self.convs.append(RGCNConv(H, O, self.num_node_types, num_edge_types))

        self.reset_parameters()

    def reset_parameters(self):
        for emb in self.emb_dict.values():
            torch.nn.init.xavier_uniform_(emb)
        for conv in self.convs:
            conv.reset_parameters()

    def group_input(self, x_dict, node_type, local_node_idx, n_id=None):
        # Create global node feature matrix.
        if n_id is not None:
            node_type = node_type[n_id]
            local_node_idx = local_node_idx[n_id]

        h = torch.zeros((node_type.size(0), self.in_channels),
                        device=node_type.device)

        for key, x in x_dict.items():
            mask = node_type == key
            h[mask] = x[local_node_idx[mask]]

        for key, emb in self.emb_dict.items():
            mask = node_type == int(key)
            h[mask] = emb[local_node_idx[mask]]

        return h

    def forward(self, n_id, x_dict, adjs, edge_type, node_type,
                local_node_idx):

        x = self.group_input(x_dict, node_type, local_node_idx, n_id)
        node_type = node_type[n_id]

        for i, (edge_index, e_id, size) in enumerate(adjs):
            x_target = x[:size[1]]  # Target node embeddings.
            node_type = node_type[:size[1]]  # Target node types.
            conv = self.convs[i]
            x = conv((x, x_target), edge_index, edge_type[e_id], node_type)
            if i != self.num_layers - 1:
                x = F.relu(x)
                x = F.dropout(x, p=0.5, training=self.training)

        return x.log_softmax(dim=-1)

    def inference(self, x_dict, edge_index_dict, key2int):
        # We can perform full-batch inference on GPU.

        device = list(x_dict.values())[0].device

        x_dict = copy(x_dict)
        for key, emb in self.emb_dict.items():
            x_dict[int(key)] = emb

        adj_t_dict = {}
        for key, (row, col) in edge_index_dict.items():
            adj_t_dict[key] = SparseTensor(row=col, col=row).to(device)

        for i, conv in enumerate(self.convs):
            out_dict = {}

            for j, x in x_dict.items():
                out_dict[j] = conv.root_lins[j](x)

            for keys, adj_t in adj_t_dict.items():
                src_key, target_key = keys[0], keys[-1]
                out = out_dict[key2int[target_key]]
                tmp = adj_t.matmul(x_dict[key2int[src_key]], reduce='mean')
                out.add_(conv.rel_lins[key2int[keys]](tmp))

            if i != self.num_layers - 1:
                for j in range(self.num_node_types):
                    F.relu_(out_dict[j])

            x_dict = out_dict

        return x_dict

    def fullBatch_inference(self, x_dict, edge_index, edge_type, node_type,
                            local_node_idx):
        x = self.group_input(x_dict, node_type, local_node_idx)
        node_type = node_type[edge_index[1].unique()]
        for i, conv in enumerate(self.convs):
            x = conv(x, edge_index, edge_type, node_type)
            if i != self.num_layers - 1:
                x = F.relu(x)
                x = F.dropout(x, p=self.dropout, training=self.training)
        return x

    def group_inference(self, x_dict, edge_index_dict, key2int):
        x_dict = copy(x_dict)
        for k, emb in self.emb_dict.items():
            x_dict[int(k)] = emb
        for i, conv in enumerate(self.convs):
            out_dict = dict()
            for n, x_target in x_dict.items():
                edge_index_n = []
                edge_type_n = []
                x_nbr = []
                for k, e_i in edge_index_dict.items():
                    if key2int[k[-1]] == n:
                        edge_index_n.append(e_i)
                        edge_type_n.append(
                            e_i.new_full((e_i.size(1), ), key2int[k]))
                        x_nbr.append(x_dict[key2int[k[0]]])
                edge_index_n = torch.cat(edge_index_n, dim=1)
                edge_type_n = torch.cat(edge_type_n, dim=0)
                x = torch.cat([x_target] + x_nbr, dim=0)
                node_type_n = edge_index_n.new_full((x_target.size(0), ), n)
                x_out = conv((x, x_target), edge_index_n, edge_type_n,
                             node_type_n)
                if i != self.num_layers - 1:
                    x_out = F.relu(x_out)
                    x_out = F.dropout(x_out,
                                      p=self.dropout,
                                      training=self.training)
                out_dict[n] = x_out
            # if i != self.num_layers - 1:
            #     for j in range(self.num_node_types):
            #         F.relu_(out_dict[j])
            x_dict = out_dict
        return x_dict
Exemple #5
0
class BRGCN(torch.nn.Module):
    def __init__(self,
                 in_channels,
                 hidden_channels,
                 out_channels,
                 num_layers,
                 dropout,
                 neg_slope,
                 heads,
                 num_relations,
                 num_nodes_dict,
                 x_types,
                 attention='Bilevel'):
        super(BRGCN, self).__init__()
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.out_channels = out_channels
        self.num_layers = num_layers
        self.dropout = dropout
        self.heads = heads
        self.num_relations = num_relations

        node_types = list(num_nodes_dict.keys())
        num_node_types = len(node_types)

        self.emb_dict = ParameterDict({
            f'{key}': Parameter(torch.Tensor(num_nodes_dict[key], in_channels))
            for key in set(node_types).difference(set(x_types))
        })

        self.convs = ModuleList()
        if attention == 'Bilevel':
            self.convs.append(
                BRGCNConv(self.in_channels, self.hidden_channels, neg_slope,
                          self.num_relations, num_node_types, self.heads,
                          self.dropout))
            for _ in range(num_layers - 2):
                self.convs.append(
                    BRGCNConv(self.hidden_channels, self.hidden_channels,
                              neg_slope, self.num_relations, num_node_types,
                              self.heads, self.dropout))
            self.convs.append(
                BRGCNConv(self.heads * self.hidden_channels, self.out_channels,
                          neg_slope, self.num_relations, num_node_types, 1,
                          self.dropout))
        elif attention == 'Node':
            self.convs = ModuleList()
            self.convs.append(
                BRGCNConvNode(self.in_channels, self.hidden_channels,
                              neg_slope, self.num_relations, num_node_types,
                              self.heads, self.dropout))
            for _ in range(num_layers - 2):
                self.convs.append(
                    BRGCNConvNode(self.hidden_channels, self.hidden_channels,
                                  self.neg_slope, self.num_relations,
                                  num_node_types, self.heads, self.dropout))
            self.convs.append(
                BRGCNConvNode(self.heads * self.hidden_channels,
                              self.out_channels, self.neg_slope,
                              self.num_relations, num_node_types, 1,
                              self.dropout))
        elif attention == 'Relation':
            self.convs = ModuleList()
            self.convs.append(
                BRGCNConvRel(self.in_channels, self.hidden_channels, neg_slope,
                             self.num_relations, num_node_types, self.heads,
                             self.dropout))
            for _ in range(num_layers - 2):
                self.convs.append(
                    BRGCNConvRel(self.hidden_channels, self.hidden_channels,
                                 self.neg_slope, self.num_relations,
                                 self.heads, num_node_types, self.dropout))
            self.convs.append(
                BRGCNConvRel(self.heads * self.hidden_channels,
                             self.out_channels, self.neg_slope,
                             self.num_relations, num_node_types, 1,
                             self.dropout))
        else:
            raise NotImplementedError

        self.reset_parameters()

    def reset_parameters(self):
        for embedding in self.emb_dict.values():
            init.xavier_uniform_(embedding)
        for conv in self.convs:
            conv.reset_parameters()

    def group_input(self, x_dict, node_type, local_node_idx, n_id=None):
        # Create global node feature matrix.
        if n_id is not None:
            node_type = node_type[n_id]
            local_node_idx = local_node_idx[n_id]

        h = torch.zeros((node_type.size(0), self.in_channels),
                        device=node_type.device)

        for key, x in x_dict.items():
            mask = node_type == key
            h[mask] = x[local_node_idx[mask]]

        for key, emb in self.emb_dict.items():
            mask = node_type == int(key)
            h[mask] = emb[local_node_idx[mask]]

        return h

    def forward(self, n_id, x_dict, adjs, edge_type, node_type,
                local_node_idx):
        '''
        :param n_id: node index for the source nodes
        :param x_dict: Node embedding dictionary with node type  as key
        :param adjs: source to node structure, (edge_index, e_id, size)
        :param edge_type: the edge type for all edges
        :param node_type: the node type for all nodes
        :param local_node_idx: transform the global node id to a local one order for each type of node
        :return:
        '''
        x = self.group_input(x_dict, node_type, local_node_idx, n_id)
        node_type = node_type[n_id]

        for i, (edge_index, e_id, size) in enumerate(adjs):
            x_target = x[:size[1]]  # Target node embeddings.
            node_type = node_type[:size[1]]  # Target node types.
            conv = self.convs[i]
            x = conv((x, x_target), edge_index, edge_type[e_id], node_type)
            if i != self.num_layers - 1:
                x = F.relu(x)
                x = F.dropout(x, p=self.dropout, training=self.training)

        return x.log_softmax(dim=-1)

    def fullBatch_inference(self, x_dict, edge_index, edge_type, node_type,
                            local_node_idx):
        x = self.group_input(x_dict, node_type, local_node_idx)
        node_type = node_type[edge_index[1].unique()]
        for i, conv in enumerate(self.convs):
            x = conv(x, edge_index, edge_type, node_type)
            if i != self.num_layers - 1:
                x = F.relu(x)
                x = F.dropout(x, p=self.dropout, training=self.training)
        return x

    def group_inference(self, x_dict, edge_index_dict, key2int):
        x_dict = copy(x_dict)
        for k, emb in self.emb_dict.items():
            x_dict[int(k)] = emb
        for i, conv in enumerate(self.convs):
            out_dict = dict()
            for n, x_target in x_dict.items():
                edge_index_n = []
                edge_type_n = []
                node_type_nbr = []
                x_nbr = []
                for k, e_i in edge_index_dict.items():
                    if key2int[k[-1]] == n:
                        edge_index_n.append(e_i)
                        edge_type_n.append(
                            e_i.new_full((e_i.size(1), ), key2int[k]))
                        node_type_nbr.append(
                            e_i.new_full((x_dict[key2int[k[0]]].size(0), ),
                                         key2int[k[0]]))
                        x_nbr.append(x_dict[key2int[k[0]]])
                edge_index_n = torch.cat(edge_index_n, dim=1)
                edge_type_n = torch.cat(edge_type_n, dim=0)
                x = torch.cat([x_target] + x_nbr, dim=0)
                node_type_target = edge_index_n.new_full((x_target.size(0), ),
                                                         n)
                node_type_n = torch.cat([node_type_target] + node_type_nbr,
                                        dim=0)
                x_out = conv((x, x_target), edge_index_n, edge_type_n,
                             node_type_n)
                if i != self.num_layers - 1:
                    x_out = F.relu(x_out)
                    x_out = F.dropout(x_out,
                                      p=self.dropout,
                                      training=self.training)
                out_dict[n] = x_out
            # if i != self.num_layers - 1:
            #     for j in range(self.num_node_types):
            #         F.relu_(out_dict[j])
            x_dict = out_dict
        return x_dict