boolean_mask o producto de puntos dispersos en tensorflow

tl; dr, ¿cuál es la forma más eficiente de elegir dinámicamente algunas entradas de un tensor?

Estoy tratando de implementar GCN sintáctica en Tensorflow. Básicamente, necesito tener una matriz de ponderación diferente para cada etiqueta (ignoremos los sesgos para esta pregunta) y elegir en cada ejecución las entradas relevantes a usar, esas serían elegidas por una matriz dispersa (para cada entrada hay como máximo una etiqueta en una dirección y en su mayoría sin bordes, ni siquiera eso).

Más concretamente, cuando tengo una matriz dispersa de bordes etiquetados (cero-uno), ¿es mejor usarla en una máscara, en una multiplicación del tensor poco denso o tal vez solo en la multiplicación normal (supongo que no en esta última, pero por simplicidad?) utilízalo en el ejemplo)

ejemplo:

units = 6 # output size x = ops.convert_to_tensor(inputs[0], dtype=self.dtype) labeled_edges = ops.convert_to_tensor(inputs[1], dtype=self.dtype) edges_shape = labeled_edges.get_shape().as_list() labeled_edges = expand_dims(labeled_edges, -2) labeled_edges = tile( labeled_edges, [1] * (len(edges_shape) - 1) + [units, 1]) graph_kernel = math_ops.multiply(self.kernel, labeled_edges) # here is the question basically outputs = standard_ops.tensordot(x, graph_kernel, [[1], [0]]) outputs = math_ops.reduce_sum(outputs, [-1]) 

Para responder a su pregunta tl; dr, puede intentar usar uno de los siguientes:

  • tf.nn.embedding_lookup : el uso típico es tf.nn.embedding_lookup(params, ids) . Devuelve un Tensor , cuyas entradas en el eje 0 son un subconjunto de parámetros de Tensor . Los índices de entradas guardadas están definidos por los identificadores de Tensor .

  • tf.nn.embedding_lookup_sparse : es lo mismo que tf.nn.embedding_lookup pero toma los ids como un SparseTensor .