Using eager execution may mean undoing some habits developed with TensorFlow graphs since code snippets that used to run once (e.g., Python function that constructs the graph to compute the loss) will run repeatedly (the same Python function will now compute the loss on each iteration).
I took a cursory look at code links provided and noticed some easy wins that would probably also be seen by using standard Python profiling tools. You may want use those (cProfile, pyspy etc.)
For example, the Keras network is currently implemented as:
class NFModel(tf.keras.Model):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def call(self, *args, **kwargs):
num_layers = 6
d, r = 2, 2
bijectors = []
for i in range(num_layers):
with tf.variable_scope('bijector_%d' % i):
V = tf.get_variable('V', [d, r], dtype=DTYPE) # factor loading
shift = tf.get_variable('shift', [d], dtype=DTYPE) # affine shift
L = tf.get_variable('L', [d * (d + 1) / 2], dtype=DTYPE) # lower triangular
bijectors.append(tfb.Affine(
scale_tril=tfd.fill_triangular(L),
scale_perturb_factor=V,
shift=shift,
))
alpha = tf.get_variable('alpha', [], dtype=DTYPE)
abs_alpha = tf.abs(alpha) + .01
bijectors.append(LeakyReLU(alpha=abs_alpha))
base_dist = tfd.MultivariateNormalDiag(loc=tf.zeros([2], DTYPE))
mlp_bijector = tfb.Chain(list(reversed(bijectors[:-1])), name='2d_mlp_bijector')
dist = tfd.TransformedDistribution(distribution=base_dist, bijector=mlp_bijector)
Instead, if you create the variables in __init__ once and avoid tf.get_variable calls on every call to the network, you should see a big improvement.
class NFModel(tf.keras.Model):
def __init__(self, *args, **kwargs):
super(NFModel, self).__init__(*args, **kwargs)
num_layers = 6
d, r = 2, 2
self.num_layers = num_layers
self.V = [tf.get_variable('V', [d, r], dtype=DTYPE) for _ in range(num_layers)]
self.shift = [tf.get_variable('shift', [d], dtype=DTYPE) for _ in range(num_layers)]
self.L = [tf.get_variable('L', [d * (d + 1) / 2], dtype=DTYPE) for _ in range(num_layers)]
self.alpha = [tf.get_variable('alpha', [], dtype=DTYPE) for _ in range(num_layers)]
def call(self, *args, **kwargs):
bijectors = []
for i in range(self.num_layers):
V = self.V[i]
shift = self.shift[i]
L = self.L[i]
bijectors.append(tfb.Affine(
scale_tril=tfd.fill_triangular(L),
scale_perturb_factor=V,
shift=shift,
))
alpha = self.alpha[i]
abs_alpha = tf.abs(alpha) + .01
bijectors.append(LeakyReLU(alpha=abs_alpha))
base_dist = tfd.MultivariateNormalDiag(loc=tf.zeros([2], DTYPE))
mlp_bijector = tfb.Chain(list(reversed(bijectors[:-1])), name='2d_mlp_bijector')
dist = tfd.TransformedDistribution(distribution=base_dist, bijector=mlp_bijector)
return {"dist": dist}
There are probably other such easy wins, a profiling tool will nudge you in the right direction.
Also, note that, TF 2.0 is less about "eager execution" and more about how one interacts with graphs, as per the RFC
Hope that helps.
