如何在TensorFlow中有效实现此简单结构

Question

小问题。如何在Python3的TensorFlow中使用密集层显式实现以下函数f（x），以实现快速构建和评估？

我尝试首先使用列表然后进行串联，但是函数f_TF的构建时间太可怕了，请参见下面的最小示例代码的输出（将近2秒的时间仅在笔记本电脑上构建128个函数，评估时间很好，最后用提取的权重W和变量符号vs中的偏差b结束时用numpy进行最后检查。我相信有一些有效的方法可以有效地构建它，但是我现在看不到它。

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

您可以创建一个形状为(N, N, M)的单个权重矩阵，然后使用tf.tensordot()为所有这些tf.matmul()矩阵一起计算W：

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

使用tf.layers中的一层来一次计算所有内容，这是一种可行的方法：

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

如the other answer中所述，您也可以只使用变量和tf.tensordot或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)