hollow dust
In the current project I am working at, we are using tf2 but tensorflow.compat.v1 as tf.
I am using tf.signal.fft2d and tf.signal.ifft2d to filter frequencies from images.
It works but the program slowly starts using more and more memory during the training and when the GPU gets full it just stops the training.
I have tried to do many different things like memory allocator growth and others. I am sure that the fft2d and ifft2d are the problem since I isolated the code from them and the problem disappeared, so its something with them and cuFFT, I believe.
I am lost in what I can do to fix this.
The function in question:
def apply_fft_filter_tensor(images: tf.Tensor, mask: tf.Tensor) -> tf.Tensor:
"""
images: float32 Tensor, shape [B, H, W, C]
freq_ranges: float32 Tensor, shape [n_ranges, 2]
Returns: real-valued float32 Tensor [B, H, W, C]
"""
# 1. Shapes
B = tf.shape(images)[0]
H = tf.shape(images)[1]
W = tf.shape(images)[2]
C = images.shape[3] # static channel count
# 3. Flatten batch & channels → [B*C, H, W]
x = tf.reshape(images, [B * C, H, W])
x_c = tf.cast(x, tf.complex64)
# 4. FFT, mask, IFFT
Xf = tf.signal.fft2d(x_c) #MEMORY GROWTH HERE
Xf_f = Xf * mask
x_ifft = tf.signal.ifft2d(Xf_f) #MEMORY GROWTH HERE
# 5. Back to real, reshape → [B, H, W, C]
x_real = tf.math.real(x_ifft)
x_out = tf.reshape(x_real, [B, H, W, C])
x_out = tf.ensure_shape(x_out, [None, None, None, C])
return x_out
and the mask used is created outside the loop once.
tender cloak