block.py

# Copyright 2018 The GraphNets Authors. All Rights Reserved.
#
# 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.
# ============================================================================
"""Model architectures for the demos in TensorFlow 2."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from graph_nets import modules
from graph_nets import utils_tf
from six.moves import range
import sonnet as snt
import tensorflow as tf


_EDGE_BLOCK_OPT = {
    "use_edges": True,
    "use_receiver_nodes": True,
    "use_sender_nodes": True,
    "use_globals": True,
    }

_NODE_BLOCK_OPT = {
    "use_received_edges": True,
    "use_sent_edges": False,
    "use_nodes": True,
    "use_globals": True,
    }

_GLOBAL_BLOCK_OPT = {
    "use_edges": False,
    "use_nodes": True,
    "use_globals": True,
    }

def make_mlp_model(latent_size, num_layers, activate_final=True):
        """Instantiates a new MLP, followed by LayerNorm.

        The parameters of each new MLP are not shared with others generated by
        this function.

        Returns:
        A Sonnet module which contains the MLP and LayerNorm.
        """
        return snt.Sequential([
            snt.nets.MLP([latent_size] * num_layers, activate_final=activate_final),
            snt.LayerNorm(axis=-1, create_offset=True, create_scale=True)
        ])


class MLPGraphIndependent(snt.Module):
    """GraphIndependent with MLP edge, node, and global models."""

    def __init__(self,
               edge_model_fn=None,
               node_model_fn=None,
               global_model_fn=None,
               name="MLPGraphIndependent"):
        super(MLPGraphIndependent, self).__init__(name=name)
        self._network = modules.GraphIndependent(
            edge_model_fn=edge_model_fn,
            node_model_fn=node_model_fn,
            global_model_fn=global_model_fn)

    def __call__(self, inputs):
        return self._network(inputs)


class MLPGraphNetwork(snt.Module):
    """GraphNetwork with MLP edge, node, and global models."""

    def __init__(self, name='MLPGraphNetwork', **model_config):
        super(MLPGraphNetwork, self).__init__(name=name)

        if model_config['reducer'] == 'mean':
            reducer=tf.math.unsorted_segment_mean
        elif model_config['reducer'] == 'sum':
            reducer=tf.math.unsorted_segment_sum

        mlp_fn = lambda: make_mlp_model(model_config['latent_size'], model_config['num_layers'])

        self._network = modules.GraphNetwork(
            edge_model_fn=mlp_fn,
            node_model_fn=mlp_fn,
            global_model_fn=mlp_fn,
            reducer=reducer,
            edge_block_opt=model_config['edge_block_opt'],
            node_block_opt=model_config['node_block_opt'],
            global_block_opt=model_config['global_block_opt'])
                                         

    def __call__(self, inputs):
        return self._network(inputs)


class MLPDeepSets(snt.Module):
    """GraphNetwork with MLP edge, node, and global models."""

    def __init__(self, name='MLPDeepSets', **model_config):
        super(MLPDeepSets, self).__init__(name=name)

        if model_config['reducer'] == 'mean':
            reducer=tf.math.unsorted_segment_mean
        elif model_config['reducer'] == 'sum':
            reducer=tf.math.unsorted_segment_sum

        mlp_fn = lambda: make_mlp_model(model_config['latent_size'], model_config['num_layers'])

        self._network = modules.DeepSets(
            node_model_fn=mlp_fn,
            global_model_fn=mlp_fn,
            reducer=reducer)
                                         

    def __call__(self, inputs):
        return self._network(inputs)


class BlockModel(snt.Module):
    """

    """

    def __init__(self,
               node_output_size=None,
               global_output_size=None,
               model_config=None,
               name="BlockModel"):
        super(BlockModel, self).__init__(name=name)

        self._num_blocks = model_config['num_blocks']
        self._concat_input = model_config['concat_input']
        self._model_config = model_config

        if self._model_config['block_type'] == 'graphnet':
            block_type = MLPGraphNetwork
        elif self._model_config['block_type'] == 'deepsets':
            block_type = MLPDeepSets


        self._core = [
                block_type(name="core_"+str(i), **self._model_config) for i in range(self._num_blocks)
                ]


        # Transforms the outputs into the appropriate shapes.
        edge_fn = None
        if node_output_size is None:
            node_fn = None
        else:
            node_fn = lambda: snt.Linear(node_output_size, name="node_output")
        if global_output_size is None:
            global_fn = None
        else:
            global_fn = lambda: snt.Linear(global_output_size, name="global_output")

        self._output_transform = modules.GraphIndependent(
            edge_fn, node_fn, global_fn, name="network_output")

    def __call__(self, input_op):
        latent = self._core[0](input_op)
        latent_all = [input_op]
        for i in range(1, self._num_blocks):
            if self._concat_input:
                core_input = utils_tf.concat([latent, latent_all[-1]], axis=1)
            else:
                core_input = latent

            latent_all.append(latent)
            latent = self._core[i](core_input)

        latent_all.append(latent)
        stacked_latent = utils_tf.concat(latent_all, axis=1)
        output = (self._output_transform(stacked_latent))
        return output