def test_DAG_gather():
"""Test invoking DAGGather."""
# TODO(rbharath): We need more documentation about why
# these numbers work.
batch_size = 10
n_graph_feat = 30
n_atom_feat = 30
n_outputs = 75
max_atoms = 50
layer_sizes = [100]
layer = layers.DAGGather(n_graph_feat=n_graph_feat,
n_outputs=n_outputs,
max_atoms=max_atoms,
layer_sizes=layer_sizes)
atom_features = np.random.rand(batch_size, n_atom_feat)
membership = np.sort(np.random.randint(0, batch_size, size=(batch_size)))
outputs = layer([atom_features, membership])
def __init__(self,
n_tasks,
max_atoms=50,
n_atom_feat=75,
n_graph_feat=30,
n_outputs=30,
layer_sizes=[100],
layer_sizes_gather=[100],
dropout=None,
mode="classification",
n_classes=2,
uncertainty=False,
batch_size=100,
**kwargs):
"""
Parameters
----------
n_tasks: int
Number of tasks.
max_atoms: int, optional
Maximum number of atoms in a molecule, should be defined based on dataset.
n_atom_feat: int, optional
Number of features per atom.
n_graph_feat: int, optional
Number of features for atom in the graph.
n_outputs: int, optional
Number of features for each molecule.
layer_sizes: list of int, optional
List of hidden layer size(s) in the propagation step:
length of this list represents the number of hidden layers,
and each element is the width of corresponding hidden layer.
layer_sizes_gather: list of int, optional
List of hidden layer size(s) in the gather step.
dropout: None or float, optional
Dropout probability, applied after each propagation step and gather step.
mode: str, optional
Either "classification" or "regression" for type of model.
n_classes: int
the number of classes to predict (only used in classification mode)
uncertainty: bool
if True, include extra outputs and loss terms to enable the uncertainty
in outputs to be predicted
"""
if mode not in ['classification', 'regression']:
raise ValueError("mode must be either 'classification' or 'regression'")
self.n_tasks = n_tasks
self.max_atoms = max_atoms
self.n_atom_feat = n_atom_feat
self.n_graph_feat = n_graph_feat
self.n_outputs = n_outputs
self.layer_sizes = layer_sizes
self.layer_sizes_gather = layer_sizes_gather
self.dropout = dropout
self.mode = mode
self.n_classes = n_classes
self.uncertainty = uncertainty
if uncertainty:
if mode != "regression":
raise ValueError("Uncertainty is only supported in regression mode")
if dropout is None or dropout == 0.0:
raise ValueError('Dropout must be included to predict uncertainty')
# Build the model.
atom_features = Input(shape=(self.n_atom_feat,))
parents = Input(shape=(self.max_atoms, self.max_atoms), dtype=tf.int32)
calculation_orders = Input(shape=(self.max_atoms,), dtype=tf.int32)
calculation_masks = Input(shape=(self.max_atoms,), dtype=tf.bool)
membership = Input(shape=tuple(), dtype=tf.int32)
n_atoms = Input(shape=tuple(), dtype=tf.int32)
dag_layer1 = layers.DAGLayer(
n_graph_feat=self.n_graph_feat,
n_atom_feat=self.n_atom_feat,
max_atoms=self.max_atoms,
layer_sizes=self.layer_sizes,
dropout=self.dropout,
batch_size=batch_size)([
atom_features, parents, calculation_orders, calculation_masks,
n_atoms
])
dag_gather = layers.DAGGather(
n_graph_feat=self.n_graph_feat,
n_outputs=self.n_outputs,
max_atoms=self.max_atoms,
layer_sizes=self.layer_sizes_gather,
dropout=self.dropout)([dag_layer1, membership])
n_tasks = self.n_tasks
if self.mode == 'classification':
n_classes = self.n_classes
logits = Reshape((n_tasks,
n_classes))(Dense(n_tasks * n_classes)(dag_gather))
output = Softmax()(logits)
outputs = [output, logits]
output_types = ['prediction', 'loss']
loss = SoftmaxCrossEntropy()
else:
output = Dense(n_tasks)(dag_gather)
if self.uncertainty:
log_var = Dense(n_tasks)(dag_gather)
var = Activation(tf.exp)(log_var)
outputs = [output, var, output, log_var]
output_types = ['prediction', 'variance', 'loss', 'loss']
def loss(outputs, labels, weights):
diff = labels[0] - outputs[0]
return tf.reduce_mean(diff * diff / tf.exp(outputs[1]) + outputs[1])
else:
outputs = [output]
output_types = ['prediction']
loss = L2Loss()
model = tf.keras.Model(
inputs=[
atom_features,
parents,
calculation_orders,
calculation_masks,
membership,
n_atoms #, dropout_switch
],
outputs=outputs)
super(DAGModel, self).__init__(
model, loss, output_types=output_types, batch_size=batch_size, **kwargs)