添加自定义图层时训练非常慢。我发现这个张量运算在 cpu 上而不是在 gpu 上运行，我不知道为啥？

Posted 2023-03-16

【中文标题】添加自定义图层时训练非常慢。我发现这个张量运算在 cpu 上而不是在 gpu 上运行，我不知道为啥？

【英文标题】：Traing is very slow when add a custom layers. I found this tensor ops run on cpu not on gpu,I don't konw why?添加自定义图层时训练非常慢。我发现这个张量运算在 cpu 上而不是在 gpu 上运行，我不知道为什么？

【发布时间】：2021-10-06 04:59:13

【问题描述】：

tensorFlow 版本（使用下面的命令）：tensorflow 2.4 Python版本：3.7.9 CUDA/cuDNN 版本：10.1

添加自定义图层时，训练速度非常慢。我发现这个张量运算是在 cpu 上而不是在 gpu 上运行的，我不知道为什么？

L2 = RuleLayer(100,420,4, 3, name='ruleLayer')(L1)

class RuleLayer(keras.layers.Layer):
    def __init__(self,batch_size,time_steps, n_input, n_memb, **kwargs):
        super(RuleLayer, self).__init__( **kwargs)
        self.ts = time_steps
        self.n = n_input
        self.m = n_memb
        self.batch_size = batch_size

    def build(self, batch_input_shape):
        #self.batch_size = batch_input_shape[0]
        # self.batch_size = tf.shape(batch_input_shape)[0]
        super(RuleLayer, self).build(batch_input_shape)  # Be sure to call this at the end
        
    def call(self, input_):
        #for d in range(1,self.n):
            
        CP_batch = []
        # a tensor object is not assignable*, so you cannot use it on the left-hand side of an assignment.
        # build a Python list of tensors, and tf.stack() them together at the end of the loop:
        for batch in range(self.batch_size):            
            CP = []
            for ts in range(0,self.ts):
                cp = input_[batch,ts,:,0]
                c_shape = [1]
                xd_shape = [self.m]
                
                for d in range(1,self.n):
                    # append shape indizes
                    c_shape.insert(0,self.m)
                    xd_shape.insert(0,1)
                    # get cartesian product for each dimension
                    #xd = tf.reshape(input_[batch,ts,:,d], (xd_shape))
                    #c = tf.reshape(cp,(c_shape))
                    #cp = tf.matmul(c , xd)                    
                    cp = tf.matmul( tf.reshape(cp,(c_shape)) , tf.reshape(input_[batch,ts,:,d], (xd_shape)))                    
                    tf.print("Is there a GPU available: "),
                    tf.print(tf.config.list_physical_devices("GPU"))
                    tf.print("Is the Tensor on GPU #0:  "),
                    tf.print(cp.device.endswith('GPU:0'))
                    
                flat_cp = tf.reshape(cp,(1, self.m**self.n))
                CP.append(flat_cp)
            CP = tf.reshape(CP,(1, self.ts,self.m**self.n))
            CP_batch.append(CP)

        return tf.reshape(tf.stack(CP_batch), (self.batch_size,self.ts, self.m**self.n))

    def compute_output_shape(self, batch_input_shape):
        if self.n == 1:
            return tf.TensorShape([self.batch_size, self.ts,self.m])
        else:
            return tf.TensorShape([self.batch_size, self.ts,self.m** self.n])
    
    def get_config(self):
        config = 
            'time_steps':  self.ts,
            'n_input': self.n,
            'n_memb': self.m,
            'batch_size':self.batch_size
        
        base_config = super(RuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

在我的电脑上花了将近 2 个小时/epoch。 纪元 1/150 是否有可用的 GPU： [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] 是 GPU #0 上的张量： 错误的 是否有可用的 GPU： [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] 是 GPU #0 上的张量： 错误的 是否有可用的 GPU： [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] 是 GPU #0 上的张量： 错误的 是否有可用的 GPU： [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] 是 GPU #0 上的张量： 错误的 是否有可用的 GPU： [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] 是 GPU #0 上的张量： 假的

【参考方案1】：

如果您希望特定操作在您选择的设备上运行，而不是自动为您选择的设备，您可以使用 tf.device 创建设备上下文，该上下文中的所有操作都将在相同的指定设备。

例子

import tensorflow as tf
tf.debugging.set_log_device_placement(True)

# Place tensors on the GPU
with tf.device('/GPU:0'):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

# Run on the GPU
c = tf.matmul(a, b)
print(c)