def message(self, x, e=None):
x_i = self.get_i(x) # Features of self
x_j = self.get_j(x) # Features of neighbours
# Features of outgoing edges are simply the edge features
e_ij = e
# Features of incoming edges j -> i are obtained by transposing the edge features.
# Since TF does not allow transposing sparse matrices with rank > 2, we instead
# re-order a tf.range(n_edges) and use it as indices to re-order the edge
# features.
# The following two instructions are the sparse equivalent of
# tf.transpose(E, perm=(1, 0, 2))
# where E has shape (N, N, S).
reorder_idx = gen_sparse_ops.sparse_reorder(
tf.stack([self.index_i, self.index_j], axis=-1),
tf.range(tf.shape(e)[0]),
(self.n_nodes, self.n_nodes),
)[1]
e_ji = tf.gather(e, reorder_idx)
# Concatenate the features and feed to first MLP
stack_ij = tf.concat([x_i, x_j, e_ij, e_ji],
axis=-1) # Shape: (n_edges, F + F + S + S)
for stack_conv in range(0, len(self.stack_models)):
stack_ij = self.stack_models[stack_conv](stack_ij)
stack_ij = self.stack_model_acts[stack_conv](stack_ij)
return stack_ij
def build(self, input_shape):
self.n_inputs = input_shape[-1]
self.out_2D_shape = [
input_shape[1] - 4 * self.kernel_radius,
input_shape[2] - 2 * self.kernel_radius
]
self.one_kernel_indices = self.MakeHexMaskIndices(self.kernel_radius)
self.var_ind_not_ordered = self.MakeVariableIndices(
self.kernel_radius, self.num_outputs, self.n_inputs)
tmp_ones = tf.ones([self.var_ind_not_ordered.shape[0]])
self.variable_indices, _ = gen_sparse_ops.sparse_reorder(
self.var_ind_not_ordered, tmp_ones, [
self.kernel_radius * 4 + 1, self.kernel_radius * 2 + 1,
self.n_inputs, self.num_outputs
])
#print(self.variable_indices.shape)
#self.hex_mask = tf.constant(
# self.MakeHexMask(self.kernel_radius).reshape([self.kernel_radius*4+1,self.kernel_radius*2+1,1,1]) ,
# dtype=tf.float32 )
mask_ = self.MakeHexMask(int((self.out_2D_shape[1] - 1) / 2))
i_cut = int(0.5 *
((mask_.shape[1] - 1) / 2 * 4 + 1 - self.out_2D_shape[0]))
mask_ = mask_[i_cut:-i_cut, :]
for coord in self.ignore_out:
# somehow, giving tuple coordinates directly doesn't work when loading model
mask_[tuple(coord)] = 0.0
mask_ = tf.constant(mask_.reshape(
[self.out_2D_shape[0], self.out_2D_shape[1], 1]),
dtype=tf.float32)
self.out_mask = mask_
self.zeros_matrix = tf.constant(tf.ones([
self.kernel_radius * 4 + 1, self.kernel_radius * 2 + 1,
self.n_inputs, self.num_outputs
]),
dtype=tf.float32)
self.sparse_weights = tf.Variable(self.winitializer(
[self.variable_indices.shape[0]]),
name="sparse_weights",
dtype=tf.float32)
self.offset = tf.Variable(tf.zeros([1, 1, 1, self.num_outputs]),
trainable=True,
name="b_kernel",
dtype=tf.float32)
if type(self.activation) == str:
self.activation = keras.activations.get(self.activation)
def add_self_loops_indices(indices, N=None):
"""
Given the indices of a square SparseTensor, adds the diagonal entries (i, i)
and returns the reordered indices.
:param indices: Tensor of rank 2, the indices to a SparseTensor.
:param N: the size of the N x N SparseTensor indexed by the indices. If `None`,
N is calculated as the maximum entry in the indices plus 1.
:return: Tensor of rank 2, the indices to a SparseTensor.
"""
N = tf.reduce_max(indices) + 1 if N is None else N
row, col = indices[..., 0], indices[..., 1]
mask = tf.ensure_shape(row != col, row.shape)
sl_indices = tf.range(N, dtype=row.dtype)[:, None]
sl_indices = tf.repeat(sl_indices, 2, -1)
indices = tf.concat((indices[mask], sl_indices), 0)
dummy_values = tf.ones_like(indices[:, 0])
indices, _ = gen_sparse_ops.sparse_reorder(indices, dummy_values, (N, N))
return indices
