Augmenter chaque élément d'un tenseur par le prédécesseur dans Tensorflow 2.0

Je vais vous donner quelques méthodes différentes pour mettre cela en œuvre. Je pense que la solution la plus évidente est d'utiliser tf.scan :

import tensorflow as tf

def apply_momentum_scan(m, p, axis=0):
    # Put axis first
    axis = tf.convert_to_tensor(axis, dtype=tf.int32)
    perm = tf.concat([[axis], tf.range(axis), tf.range(axis + 1, tf.rank(m))], axis=0)
    m_t = tf.transpose(m, perm)
    # Do computation
    res_t = tf.scan(lambda a, x: a * p + x, m_t)
    # Undo transpose
    perm_t = tf.concat([tf.range(1, axis + 1), [0], tf.range(axis + 1, tf.rank(m))], axis=0)
    return tf.transpose(res_t, perm_t)

Cependant, vous pouvez également l'implémenter en tant que produit matriciel particulier, si vous construisez une matrice de facteurs exponentiels :

import tensorflow as tf

def apply_momentum_matmul(m, p, axis=0):
    # Put axis first and reshape
    m = tf.convert_to_tensor(m)
    p = tf.convert_to_tensor(p)
    axis = tf.convert_to_tensor(axis, dtype=tf.int32)
    perm = tf.concat([[axis], tf.range(axis), tf.range(axis + 1, tf.rank(m))], axis=0)
    m_t = tf.transpose(m, perm)
    shape_t = tf.shape(m_t)
    m_tr = tf.reshape(m_t, [shape_t[0], -1])
    # Build factors matrix
    r = tf.range(tf.shape(m_tr)[0])
    p_tr = tf.linalg.band_part(p ** tf.dtypes.cast(tf.expand_dims(r, 1) - r, p.dtype), -1, 0)
    # Do computation
    res_tr = p_tr @ m_tr
    # Reshape back and undo transpose
    res_t = tf.reshape(res_tr, shape_t)
    perm_t = tf.concat([tf.range(1, axis + 1), [0], tf.range(axis + 1, tf.rank(m))], axis=0)
    return tf.transpose(res_t, perm_t)

Cela peut également être réécrit pour éviter la première transposition (qui dans TensorFlow est coûteuse) avec tf.tensordot :

import tensorflow as tf

def apply_momentum_tensordot(m, p, axis=0):
    # Put axis first and reshape
    m = tf.convert_to_tensor(m)
    # Build factors matrix
    r = tf.range(tf.shape(m)[axis])
    p_mat = tf.linalg.band_part(p ** tf.dtypes.cast(tf.expand_dims(r, 1) - r, p.dtype), -1, 0)
    # Do computation
    res_t = tf.linalg.tensordot(m, p_mat, axes=[[axis], [1]])
    # Transpose
    last_dim = tf.rank(res_t) - 1
    perm_t = tf.concat([tf.range(axis), [last_dim], tf.range(axis, last_dim)], axis=0)
    return tf.transpose(res_t, perm_t)

Les trois fonctions seraient utilisées de manière similaire :

import tensorflow as tf

p = tf.Variable(0.5, dtype=tf.float32)
m = tf.constant([[0, 1, 2, 3, 4],
                 [1, 3, 5, 7, 10],
                 [1, 1, 1, -1, 0]], tf.float32)
# apply_momentum is one of the functions above
print(apply_momentum(m, p, axis=0).numpy())
# [[ 0.    1.    2.    3.    4.  ]
#  [ 1.    3.5   6.    8.5  12.  ]
#  [ 1.5   2.75  4.    3.25  6.  ]]
print(apply_momentum(m, p, axis=1).numpy())
# [[ 0.      1.      2.5     4.25    6.125 ]
#  [ 1.      3.5     6.75   10.375  15.1875]
#  [ 1.      1.5     1.75   -0.125  -0.0625]]

L'utilisation d'un produit matriciel est asymptotiquement plus complexe, mais elle peut être plus rapide que la numérisation. Voici un petit benchmark :

import tensorflow as tf
import numpy as np

# Make test data
tf.random.set_seed(0)
p = tf.constant(0.5, dtype=tf.float32)
m = tf.random.uniform([100, 30, 50], dtype=tf.float32)

# Axis 0
print(np.allclose(apply_momentum_scan(m, p, 0).numpy(), apply_momentum_matmul(m, p, 0).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 0).numpy(), apply_momentum_tensordot(m, p, 0).numpy()))
# True
%timeit apply_momentum_scan(m, p, 0)
# 11.5 ms ± 610 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit apply_momentum_matmul(m, p, 0)
# 1.36 ms ± 18.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit apply_momentum_tensordot(m, p, 0)
# 1.62 ms ± 7.39 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# Axis 1
print(np.allclose(apply_momentum_scan(m, p, 1).numpy(), apply_momentum_matmul(m, p, 1).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 1).numpy(), apply_momentum_tensordot(m, p, 1).numpy()))
# True
%timeit apply_momentum_scan(m, p, 1)
# 4.27 ms ± 60.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit apply_momentum_matmul(m, p, 1)
# 1.27 ms ± 36.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit apply_momentum_tensordot(m, p, 1)
# 1.2 ms ± 11.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# Axis 2
print(np.allclose(apply_momentum_scan(m, p, 2).numpy(), apply_momentum_matmul(m, p, 2).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 2).numpy(), apply_momentum_tensordot(m, p, 2).numpy()))
# True
%timeit apply_momentum_scan(m, p, 2)
# 6.29 ms ± 64.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit apply_momentum_matmul(m, p, 2)
# 1.41 ms ± 21.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit apply_momentum_tensordot(m, p, 2)
# 1.05 ms ± 26 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Ainsi, le produit matriciel semble gagner. Voyons si cela évolue :

import tensorflow as tf
import numpy as np

# Make test data
tf.random.set_seed(0)
p = tf.constant(0.5, dtype=tf.float32)
m = tf.random.uniform([1000, 300, 500], dtype=tf.float32)

# Axis 0
print(np.allclose(apply_momentum_scan(m, p, 0).numpy(), apply_momentum_matmul(m, p, 0).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 0).numpy(), apply_momentum_tensordot(m, p, 0).numpy()))
# True
%timeit apply_momentum_scan(m, p, 0)
# 784 ms ± 6.78 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_matmul(m, p, 0)
# 1.13 s ± 76.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_tensordot(m, p, 0)
# 1.3 s ± 27 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Axis 1
print(np.allclose(apply_momentum_scan(m, p, 1).numpy(), apply_momentum_matmul(m, p, 1).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 1).numpy(), apply_momentum_tensordot(m, p, 1).numpy()))
# True
%timeit apply_momentum_scan(m, p, 1)
# 852 ms ± 12.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_matmul(m, p, 1)
# 659 ms ± 10.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_tensordot(m, p, 1)
# 741 ms ± 19.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Axis 2
print(np.allclose(apply_momentum_scan(m, p, 2).numpy(), apply_momentum_matmul(m, p, 2).numpy()))
# True
print(np.allclose(apply_momentum_scan(m, p, 2).numpy(), apply_momentum_tensordot(m, p, 2).numpy()))
# True
%timeit apply_momentum_scan(m, p, 2)
# 1.06 s ± 16.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_matmul(m, p, 2)
# 924 ms ± 17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit apply_momentum_tensordot(m, p, 2)
# 483 ms ± 10.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Eh bien, maintenant ce n'est plus aussi clair. La numérisation n'est toujours pas super rapide, mais les produits matriciels sont parfois plus lents. Comme vous pouvez l'imaginer, si vous optez pour des tenseurs encore plus grands, la complexité des produits matriciels dominera les délais.

Donc, si vous voulez la solution la plus rapide et que vous savez que vos tenseurs ne vont pas devenir énormes, utilisez l'une des implémentations de produits matriciels. Si vous êtes d'accord avec une vitesse correcte mais que vous voulez vous assurer que vous ne manquez pas de mémoire (la solution matricielle prend également beaucoup plus) et que le timing est prévisible, vous pouvez utiliser la solution d'analyse.

Remarque :Les analyses comparatives ci-dessus ont été effectuées sur le processeur, les résultats peuvent varier considérablement sur le processeur graphique.

Voici une réponse qui fournit simplement quelques informations et une solution naïve pour corriger le code --- pas le problème réel (veuillez vous référer ci-dessous pour le pourquoi).

Tout d'abord, le TypeError est un problème de types incompatibles dans les tenseurs de votre première tentative. Certains tenseurs contiennent des nombres à virgule flottante (double), certains contiennent des entiers. Cela aurait aidé à montrer le plein message d'erreur :

TypeError: Input 'y' of 'Mul' Op has type int32 that does not match type float64 of argument 'x'.

Ce qui arrive à mettre sur la bonne voie (malgré les détails sanglants de la trace de la pile).

Voici une solution naïve pour faire fonctionner le code (avec des mises en garde contre le problème cible) :

import tensorflow as tf

@tf.function
def vectorize_predec(t, p):
    _p = tf.transpose(
        tf.convert_to_tensor(
            [p * t[...,idx] for idx in range(t.shape[-1] - 1)],
            dtype=tf.float64))
    _p = tf.concat([
        tf.zeroes((_p.shape[0], 1), dtype=tf.float64),
        _p
    ], axis=1)
    return t + _p

p = tf.Variable(0.5, dtype='double')

m = tf.constant([[0, 1, 2, 3, 4],
          [1, 3, 5, 7, 10],
          [1, 1, 1, -1, 0]], dtype=tf.float64)

n = tf.constant([[0.0, 1.0, 2.5, 4.0, 5.5],
          [1.0, 3.5, 6.5, 9.5, 13.5],
          [1.0, 1.5, 1.5, -0.5, -0.5]], dtype=tf.float64)
print(f'Expected: {n}')

result = vectorize_predec(m, p)
print(f'Result: {result}')

tf.test.TestCase().assertAllEqual(n, result)

Les principaux changements :

Le m le tenseur obtient un dtype=tf.float64 pour correspondre à l'original double , donc l'erreur de type disparaît.
La fonction est essentiellement une réécriture complète. L'idée naïve est d'exploiter la définition du problème, ce qui pas indiquer si les valeurs dans N sont calculés avant ou après les mises à jour. Voici une version avant mise à jour, bien plus simple. Résoudre ce qui semble être le "vrai" problème nécessite de travailler un peu plus sur la fonction (voir d'autres réponses, et je travaillerai peut-être plus ici).

Fonctionnement de la fonction :

Il calcule les incréments attendus p * x1 , p * x2 , etc. dans un tableau Python standard. Notez qu'il s'arrête avant le dernier élément de la dernière dimension, car nous allons décaler le tableau.
Il convertit le tableau en un tenseur avec tf.convert_to_tensor , ajoutant ainsi le tableau au graphe de calcul. La transposition est nécessaire pour correspondre à la forme originale du tenseur (nous pourrions l'éviter).
Il ajoute des zéros au début de chaque dimension le long du dernier axe.
Le résultat est la somme du tenseur d'origine et du construit.

Les valeurs deviennent x1 + 0.0 * p , puis x2 + x1 * p , etc. Cela illustre quelques fonctions et problèmes à examiner (types, formes), mais j'admets que cela triche et ne résout pas le problème réel.

De plus, ce code n'est efficace sur aucun matériel. Il est juste illustratif et nécessiterait (1) d'éliminer le tableau Python, (2) d'éliminer la transposition, (3) d'éliminer l'opération de concaténation. J'espère une bonne formation :-)

Remarques supplémentaires :

Le problème demande une solution sur des tenseurs de forme (a, b, c). Le code que vous partagez fonctionne sur des tenseurs de forme (a, b), donc corriger le code ne résoudra toujours pas le problème.
Le problème nécessite des nombres rationnels. Je ne sais pas quelle est l'intention, et cette réponse laisse cette exigence de côté.
La forme de T = [x1, x2, x3, x4] est en fait (4,) , en supposant xi sont des scalaires.
Pourquoi tf.float64 ? Par défaut, nous obtenons tf.float32 , et en supprimant le double ferait fonctionner le code. Mais l'exemple perdrait le point que les types comptent, donc le choix d'un type explicite non par défaut (et d'un code plus laid).