如何使用 Keras 加速 GPU 中的函数进行模型训练答案

【问题标题】：How to speed up a function in GPU for model training using Keras如何使用 Keras 加速 GPU 中的函数进行模型训练
【发布时间】：2021-09-07 06:50:55
【问题描述】：

我在使用 Keras 训练模型之前对每个图像应用预处理功能，问题是这个复杂的功能需要时间，即使在 Google Colab Pro（基于 GPU）上，当我不使用这个功能时训练很快，所以现在，我需要一种方法来帮助我在不改变训练速度的情况下使用这个函数，我尝试了 Numba 库，但它对我不起作用，因为我的函数调用了许多其他库。

#@cuda.jit
def my_func(image):
  image = image.astype('uint8')
  image = n.transform(image)
  return image.astype('float64') 

  data_generator = ImageDataGenerator(rescale=1.0/255.0,vertical_flip=True, 
                                   horizontal_flip=True,
                                    preprocessing_function=my_func)

【问题讨论】：

不可能使用 Numba 在 GPU 上运行类似的代码。
那么，我该怎么做呢？
做什么？我已经解释过你不能在你的 GPU 上运行任意 python 代码（尤其是使用像 keras 这样的编译库的代码）。正如他们所说，没有免费的午餐。

标签： python keras gpu conv-neural-network

【解决方案1】：

尝试创建一个层来处理网络内的输入数据。

class customFn(layers.Layer):
    def __init__(self):
        super(customFn, self).__init__()

    def call(self, inputs):
        result=tf.math.divide(tf.math.divide(tf.math.square(tf.multiply(inputs,255)),255),255)
        #you should rewrite your original function based on tensorflow functions
        #like tf.math.divide
        return result

def giveModel():
    with tf.device('/device:GPU:0'):
        inputs=layers.Input(shape=inputShape)
        x=customFn()(inputs) #use it as usual
        x=layers.Conv2D(32,(3,3),(1,1),activation='relu')(x)
        #and story goes on ...

【讨论】：