Source code for stellargraph.layer.watch_your_step

# -*- coding: utf-8 -*-
#
# Copyright 2020 Data61, CSIRO
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import tensorflow as tf
from tensorflow.keras.layers import Layer, Embedding, Input, Lambda, Concatenate, Dense
from tensorflow.keras import backend as K
import numpy as np
import warnings
from ..mapper.adjacency_generators import AdjacencyPowerGenerator
from ..core.experimental import experimental


class AttentiveWalk(Layer):
    """
    This implements the graph attention as in Watch Your Step: Learning Node Embeddings via Graph Attention
    https://arxiv.org/pdf/1710.09599.pdf.

    Args:
        walk_length (int): the length of the random walks. Equivalent to the number of adjacency powers used. Defaults
            to `10` as this value was found to perform well by the authors of the paper.
        attention_initializer (str or func, optional): The initialiser to use for the attention weights.
        attention_regularizer (str or func, optional): The regulariser to use for the attention weights.
        attention_constraint (str or func, optional): The constraint to use for the attention weights.
        input_dim (tuple of ints, optional): The shape of the input to the layer.
    """

    def __init__(
        self,
        walk_length=10,
        attention_initializer="glorot_uniform",
        attention_regularizer=None,
        attention_constraint=None,
        input_dim=None,
        **kwargs,
    ):

        if "input_shape" not in kwargs and input_dim is not None:
            kwargs["input_shape"] = input_dim

        self.walk_length = walk_length
        self.attention_initializer = attention_initializer
        self.attention_regularizer = attention_regularizer
        self.attention_constraint = attention_constraint
        super().__init__(**kwargs)

    def compute_output_shape(self, input_shapes):
        return (input_shapes[0][-1],)

    compute_output_shape.__doc__ = Layer.compute_output_shape.__doc__

    def build(self, input_shapes):

        self.attention_weights = self.add_weight(
            shape=(self.walk_length,),
            initializer=self.attention_initializer,
            name="attention_weights",
            regularizer=self.attention_regularizer,
            constraint=self.attention_constraint,
        )

        self.built = True

    build.__doc__ = Layer.build.__doc__

    def call(self, partial_powers):
        """
        Applies the layer and calculates the expected random walks.

        Args:
            partial_powers: num_rows rows of the first num_powers powers of adjacency matrix with shape
            (num_rows, num_powers, num_nodes)

        Returns:
            Tensor that represents the expected random walks starting from nodes corresponding to the input rows of
            shape (num_rows, num_nodes)
        """

        attention = K.softmax(self.attention_weights)
        expected_walk = tf.einsum("ijk,j->ik", partial_powers, attention)

        return expected_walk


@experimental(reason="lack of unit tests", issues=[804])
def get_embeddings(model):
    """
    This function returns the embeddings from a model with Watch Your Step embeddings.

    Args:
        model (keras Model): a keras model that contains Watch Your Step embeddings.

    Returns:
        embeddings (np.array): a numpy array of the model's embeddings.
    """
    embeddings = np.hstack(
        [
            model.get_layer("WATCH_YOUR_STEP_LEFT_EMBEDDINGS").embeddings.numpy(),
            model.get_layer("WATCH_YOUR_STEP_RIGHT_EMBEDDINGS")
            .kernel.numpy()
            .transpose(),
        ]
    )

    return embeddings


[docs]@experimental(reason="lack of unit tests", issues=[804]) class WatchYourStep: """ Implementation of the node embeddings as in Watch Your Step: Learning Node Embeddings via Graph Attention https://arxiv.org/pdf/1710.09599.pdf. This model requires specification of the number of random walks starting from each node, and the embedding dimension to use for the node embeddings. Args: generator (AdjacencyPowerGenerator): the generator num_walks (int): the number of random walks starting at each node to use when calculating the expected random walks. Defaults to `80` as this value was found to perform well by the authors of the paper. embedding dimension (int): the dimension to use for the node embeddings (must be an even number). attention_initializer (str or func, optional): The initialiser to use for the attention weights. attention_regularizer (str or func, optional): The regulariser to use for the attention weights. attention_constraint (str or func, optional): The constraint to use for the attention weights. """ def __init__( self, generator, num_walks=80, embedding_dimension=64, attention_initializer="glorot_uniform", attention_regularizer=None, attention_constraint=None, ): if not isinstance(generator, AdjacencyPowerGenerator): raise TypeError( "generator should be an instance of AdjacencyPowerGenerator." ) if not isinstance(num_walks, int): raise TypeError("num_walks should be an int.") if num_walks <= 0: raise ValueError("num_walks should be a positive int.") self.num_walks = num_walks self.num_powers = generator.num_powers self.n_nodes = int(generator.Aadj_T.shape[0]) if embedding_dimension % 2 != 0: warnings.warn( f"embedding_dimension: expected even number, found odd number ({embedding_dimension}). It will be rounded down to {embedding_dimension - 1}." ) embedding_dimension -= 1 self.embedding_dimension = embedding_dimension self.attention_regularizer = attention_regularizer self.attention_initializer = attention_initializer self.attention_constraint = attention_constraint
[docs] def build(self): """ This function builds the layers for a keras model. returns: A tuple of (inputs, outputs) to use with a keras model. """ input_rows = Input(batch_shape=(None,), name="row_node_ids", dtype="int64") input_powers = Input(batch_shape=(None, self.num_powers, self.n_nodes)) left_embedding = Embedding( self.n_nodes, int(self.embedding_dimension / 2), input_length=None, name="WATCH_YOUR_STEP_LEFT_EMBEDDINGS", ) vectors_left = left_embedding(input_rows) # all right embeddings are used in every batch. to avoid unnecessary lookups the right embeddings are stored # in a dense layer to enable efficient dot product between the left vectors in the current batch and all right # vectors outer_product = Dense( self.n_nodes, use_bias=False, kernel_initializer="uniform", name="WATCH_YOUR_STEP_RIGHT_EMBEDDINGS", )(vectors_left) sigmoids = tf.keras.activations.sigmoid(outer_product) attentive_walk_layer = AttentiveWalk( walk_length=self.num_powers, attention_constraint=self.attention_constraint, attention_regularizer=self.attention_regularizer, attention_initializer=self.attention_initializer, ) expected_walk = self.num_walks * attentive_walk_layer(input_powers) # layer to add batch dimension of 1 to output expander = Lambda(lambda x: K.expand_dims(x, axis=1)) output = Concatenate(axis=1)([expander(expected_walk), expander(sigmoids)]) return [input_rows, input_powers], output