feature transform running issue on GPU
ehsanngh opened this issue · comments
Hello,
In my code, I'm using a Fourier transform like the one you used in "holography." The problem is that my code becomes very very slow on GPU while it can be executed on CPU with an acceptable runtime. I'm reaching out to see whether you faced the same issue or not.
def feature_transform(inputs):
# Periodic BC in x
w = 2 * np.pi / P
x, y = w * inputs[:, :1], inputs[:, 1:]
return tf.concat(
(
tf.math.cos(x),
tf.math.sin(x),
tf.math.cos(2 * x),
tf.math.sin(2 * x),
tf.math.cos(3 * x),
tf.math.sin(3 * x),
tf.math.cos(4 * x),
tf.math.sin(4 * x),
tf.math.cos(5 * x),
tf.math.sin(5 * x),
tf.math.cos(6 * x),
tf.math.sin(6 * x),
y,
tf.math.cos(OMEGA * y),
tf.math.sin(OMEGA * y),
),
axis=1,
)
To explain more about this issue, I found that the tensor slicing like this inputs[:, :1]
and the function tf.concat
are making the code slow. I just want to know whether you faced this issue before. Thanks for your time.
Update: Switching the backend from tf2
to tf1
will solve the issue.