keras 模型上奇怪的分析结果：越复杂越快

Posted 2023-03-06

【中文标题】keras 模型上奇怪的分析结果：越复杂越快

【英文标题】：Weird profiling result on keras models: the more complex the faster

【发布时间】：2019-08-28 15:43:08

【问题描述】：

所以我目前正在尝试找出最适合处理傅立叶变换的深度学习框架。到目前为止，我使用 keras 和 tensorflow 后端，但我注意到 fft 有点慢（参见 this issue on Github）。

所以最近我尝试直接和pytorch比较速度。由于我想做的不仅仅是简单地进行傅立叶变换，因此我尝试添加一些操作来进行更全面的基准测试，我注意到对于keras，添加操作正在减少计算时间。

这是最小的工作示例（基本上是在 2D 中进行逆傅立叶变换，通过获取图像的模块完成，并且介于可能的“去复杂化”和“重新复杂化”之间）：

import os
os.environ["CUDA_VISIBLE_DEVICES"]="-1"
from keras.layers import Input, Lambda, concatenate
from keras.models import Model
import numpy as np
import tensorflow as tf
from tensorflow.signal import ifft2d
import time

def concatenate_real_imag(x):
    x_real = Lambda(tf.math.real)(x)
    x_imag = Lambda(tf.math.imag)(x)
    return concatenate([x_real, x_imag])

def to_complex(x):
    return tf.complex(x[0], x[1])

def complex_from_half(x, n, output_shape):
    return Lambda(lambda x: to_complex([x[..., :n], x[..., n:]]), output_shape=output_shape)(x)

def weird_model(conc_then_com=False):
    input_size = (320, None, 1)
    kspace_input = Input(input_size, dtype='complex64', name='kspace_input')
    inv_kspace = Lambda(ifft2d, output_shape=input_size)(kspace_input)
    if conc_then_com:
        inv_kspace = concatenate_real_imag(kspace_input)
        inv_kspace = complex_from_half(inv_kspace, 1, input_size)
    abs_inv_kspace = Lambda(tf.math.abs)(inv_kspace)
    model = Model(inputs=kspace_input, outputs=abs_inv_kspace)
    model.compile(
        optimizer='adam',
        loss='mse',
    )
    return model

# fake data
data_x = np.random.rand(35, 320, 320, 1) + 1j * np.random.rand(35, 320, 320, 1)
data_y = np.random.rand(35, 320, 320, 1)

start = time.time()
r = weird_model(conc_then_com=True).predict_on_batch(data_x)
end = time.time()
duration = end - start
print(f'For the prediction with the complex model it took duration')

start = time.time()
r = weird_model(conc_then_com=False).predict_on_batch(data_x)
end = time.time()
duration = end - start
print(f'For the prediction with the simple model it took duration')

start = time.time()
weird_model(conc_then_com=True).fit(
    x=data_x,
    y=data_y,
    batch_size=35,
    epochs=1,
    verbose=2,
    shuffle=False,
)
end = time.time()
duration = end - start
print(f'For the fitting with the complex model it took duration')

start = time.time()
weird_model(conc_then_com=False).fit(
    x=data_x,
    y=data_y,
    batch_size=35,
    epochs=1,
    verbose=2,
    shuffle=False,
)
end = time.time()
duration = end - start
print(f'For the fitting with the simple model it took duration')

给出以下时间（或多或少）：

For the prediction with the complex model it took 0.24
For the prediction with the simple model it took 3.98
For the fitting with the complex model it took 0.28
For the fitting with the simple model it took 4.01

我不知道发生了什么。

【问题讨论】：

顺便说一句，如果有人有兴趣知道为什么 fft 在 tensorflow 中如此缓慢，有一个 Github issue 已关闭（但确实是开放的）来解决该主题。

【参考方案1】：

实际上，这只是一个错字： inv_kspace = concatenate_real_imag(kspace_input) 应该是 inv_kspace = concatenate_real_imag(inv_kspace)