How to efficiently implement this simple structure in TensorFlow

Question

Small question. How would you implement the following function f(x) explicitly with dense layers for a fast build time and evaluation in TensorFlow in Python3?

I tried to do this first with a list and then concatenation but the building time for the function f_TF is horrible, see output of minimal example code below (almost 2 seconds just building for 128 function on my laptop, evaluation time is fine, final check with numpy at the end with extracted weights W and biases b in variables symbols vs ). I am sure there is some obvious way to build this efficiently, but I can not see it right now.

import numpy as np
import tensorflow as tf
import datetime

#%% Parameters

# Dimension of input x number n_p of data points
n = 4
n_p = 10

# Number of functions
m = 128

#%% Generate data

# Input data x
D = np.random.rand(n_p,n)

#%% TF

# Placeholder for input x
x_p = tf.placeholder(dtype=tf.float64,shape=[None,n])

# Build f(x)
t1 = datetime.datetime.now()
fs = []
for _ in range(m):
    floc = tf.layers.dense(x_p,n,tf.square,bias_initializer=tf.glorot_normal_initializer())
    floc = tf.sin(tf.reduce_sum(floc,axis=1,keepdims=True))
    fs.append(floc)
f_TF = tf.concat(fs,axis=1)
t2 = datetime.datetime.now()
print('Time to build f(x): \n\t%s' % (t2-t1))

# Session and evaluate
sess = tf.Session()
sess.run(tf.global_variables_initializer())
t1 = datetime.datetime.now()
f_TF_values = sess.run(f_TF,{x_p:D})
t2 = datetime.datetime.now()
print('Time for TF evaluation: \n\t%s' % (t2-t1))

# Extract weights and biases
t1 = datetime.datetime.now()
vs = tf.global_variables()
vs = [sess.run(v) for v in vs]
t2 = datetime.datetime.now()
print('Time for extraction of variables: \n\t%s' % (t2-t1))

sess.close()
tf.reset_default_graph()

#%% NP

# Check single evaluation
i = np.random.randint(0,n_p)
x0 = D[i]
f0 = np.array([np.sin(np.linalg.norm(np.matmul(x0,vs[2*i])+vs[2*i+1])**2) for i in range(len(vs)//2)])
print('Deviation from single evaluation: \n\t%.2e' % np.linalg.norm(f0-f_TF_values[i]))

# Check all
t1 = datetime.datetime.now()
f_NP_values = np.hstack([
        np.sin(np.linalg.norm(np.matmul(D,vs[2*i])+vs[2*i+1],axis=1,keepdims=True)**2) 
        for i in range(len(vs)//2)])
t2 = datetime.datetime.now()
print('Time for NP evaluation: \n\t%s' % (t2-t1))

print('Deviation between TF and NP computations: \n\t%.2e' % np.linalg.norm(f_TF_values - f_NP_values))

Answer 1

You could create a single weight matrix of shape (N, N, M) and then use tf.tensordot() to compute tf.matmul() for all these W matrices together:

import tensorflow as tf
import numpy as np

N = 4
M = 3
n_samples = 5

def sin_layer(x, units):
    N = x.get_shape().as_list()[-1]
    w = tf.Variable(tf.random.normal((N, N, units)), tf.float32)
    b = tf.Variable(tf.zeros((N, units)))
    tensor = tf.tensordot(x, w, axes=[[1], [0]]) + b # <-- matmul all `W`s at once
    tensor = tf.reduce_sum(tf.square(tensor), axis=1) # <-- reduce `N` dimension
    tensor = tf.math.sin(tensor)
    return tensor

x = tf.placeholder(tf.float32, shape=(None, N))
tensor = sin_layer(x, units=M)

x_data = np.random.normal(size=(n_samples, N)) # <-- 5 samples of size `N==4`

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    res = sess.run(tensor, feed_dict={x:x_data})
    print(res.shape) # <-- `(5, 3)==(n_samples, M)`
    print(res)
# [[ 0.24542944 -0.25523183  0.9970544 ]
#  [ 0.992266   -0.98576933 -0.65339005]
#  [ 0.95481074  0.8390483   0.41041443]
#  [-0.6582102  -0.98120177 -0.00824349]
#  [ 0.61224973  0.7946086   0.6564668 ]]

Answer 2

This one possible way you can do that, using one layer from tf.layers to compute everything at once:

import tensorflow as tf
import numpy as np

# Dimension of input x number n_p of data points
n = 4
n_p = 10

# Number of functions
m = 128

#%% Generate data
np.random.seed(0)
D = np.random.rand(n_p, n)

# Placeholder for input x
x_p = tf.placeholder(dtype=tf.float64, shape=[None, n])
# Make a single big layer
layer = tf.layers.Dense(n * m, tf.square, bias_initializer=tf.glorot_normal_initializer())
floc = tf.reshape(layer(x_p), [-1, m, n])
f_TF = tf.sin(tf.reduce_sum(floc, axis=2))
# You can still retrieve the individual matrices and biases if you want
w_all, b_all = layer.weights
w = [w_all[:, i * n:(i + 1) * n] for i in range(m)]
b = [b_all[i * n:(i + 1) * n] for i in range(m)]
# Check the operation is equivalent
f_TF1 = tf.stack([tf.sin(tf.reduce_sum(tf.square(x_p @ wi + bi), axis=1)) for wi, bi in zip(w, b)], axis=1)
init = tf.global_variables_initializer()
with tf.Session() as sess:
    sess.run(init)
    v1, v2 = sess.run([f_TF, f_TF1], feed_dict={x_p: D})
    print(np.allclose(v1, v2))
    # True

As mentioned in the other answer , you can also just use variables and tf.tensordot or tf.einsum :

import tensorflow as tf
import numpy as np

# Dimension of input x number n_p of data points
n = 4
n_p = 10

# Number of functions
m = 128

#%% Generate data
np.random.seed(0)
D = np.random.rand(n_p, n)

# Placeholder for input x
x_p = tf.placeholder(dtype=tf.float64, shape=[None, n])
# Make a single big layer
w_all = tf.get_variable('W', [n, n, m], dtype=x_p.dtype)
b_all = tf.get_variable('B', [n, m], dtype=x_p.dtype, initializer=tf.glorot_normal_initializer())
floc = tf.square(tf.einsum('bi,ijm->bjm', x_p, w_all) + b_all)
f_TF = tf.sin(tf.reduce_sum(floc, axis=1))
# Matrices and biases
w = tf.unstack(w_all, axis=2)
b = tf.unstack(b_all, axis=1)