非常感谢您阅读我的问题。我对TensorFlow还比较陌生，我编写了这个简单的自定义层，以便我可以实现数据依赖的条件，根据输入选择和训练不同的变量集。
class Custom_Layer1234(keras.layers.Layer): def __init__(self, units=45, input_dim=45): super(Custom_Layer1234, self).__init__() w_init = tf.random_normal_initializer() b_init = tf.zeros_initializer() self.w_0 = tf.Variable(initial_value=w_init(shape=(input_dim, units,)),trainable=True) self.w_1 = tf.Variable(initial_value=w_init(shape=(input_dim, units,)),trainable=True) self.w_2 = tf.Variable(initial_value=w_init(shape=(input_dim, units,)),trainable=True) self.b_0 = tf.Variable(initial_value=b_init(shape=(units,)),trainable=True) self.b_1 = tf.Variable(initial_value=b_init(shape=(units,)),trainable=True) self.b_2 = tf.Variable(initial_value=b_init(shape=(units,)),trainable=True) @tf.function def call(self, inputs): my_list = tf.TensorArray(tf.float32, size=64, dynamic_size=False) //batch_size = 64 i = 0 for x in inputs: This_diff = x[1] x = tf.reshape(x,shape=(1,-1)) //if i dont reshape cannot tf.matmul if This_diff > 0 : my_list = my_list.write(i, tf.nn.relu(tf.matmul(inputs, self.w_0) + self.b_0)) elif This_diff < 0 : my_list = my_list.write(i, tf.nn.relu(tf.matmul(inputs, self.w_1) + self.b_1)) else: my_list = my_list.write(i, tf.nn.relu(tf.matmul(inputs, self.w_2) + self.b_2)) i += 1 return my_list.concat()代码运行没有问题，只是速度极慢，以至于CPU训练甚至比GPU训练还要快，但仍然需要很长时间。作为替代方案，我尝试了以下方法：
class Custom_Layer1234(keras.layers.Layer): def __init__(self, units=45, input_dim=45): ..... def Mapping_this_Function(self, x): This_diff = x[1] x = tf.reshape(x,shape=(1,-1)) if This_diff > 0 : return tf.nn.relu(tf.matmul(inputs, self.w_0) + self.b_0) elif This_diff < 0 : return tf.nn.relu(tf.matmul(inputs, self.w_1) + self.b_1) else: return tf.nn.relu(tf.matmul(inputs, self.w_2) + self.b_2) @tf.function def call(self, inputs): return tf.map_fn(self.Mapping_this_Function,inputs, parallel_iterations=64)这种方法也可以编译和拟合，但非常奇怪的是，准确率变得异常低，为1.6342e-04，与第一轮训练的通常值0.1-0.2相比，而且随时间不增反降，是否map_fn不应该在自定义层中使用。在这两种情况下，我都尝试了有无以下设置：
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"我这样处理‘批量大小张量上的for循环条件’是否正确，请您指出我可能遗漏的更有效的方法。我对代码在底层如何运行的理解非常浅薄，但发现CPU通常更适合处理顺序操作，而GPU更适合并行操作。我将输入张量切割成更小的片段，并多次使用这个自定义层，是否值得尝试用C++编写自定义操作。
非常感谢
回答：
我认为您可以像这样将所有操作向量化：
@tf.functiondef call(self, inputs): gt = tf.expand_dims(tf.cast(inputs[:, 1] > 0, tf.float32), -1) lt = tf.expand_dims(tf.cast(inputs[:, 1] < 0, tf.float32), -1) eq = tf.expand_dims(tf.cast(inputs[:, 1] == 0, tf.float32), -1) x0 = gt * tf.nn.relu(tf.matmul(inputs, self.w_0) + self.b_0) x1 = lt * tf.nn.relu(tf.matmul(inputs, self.w_1) + self.b_1) x2 = eq * tf.nn.relu(tf.matmul(inputs, self.w_2) + self.b_2) return x0 + x1 + x2由于这三个条件是互斥的，对于每个i，{gt[i], lt[i], eq[i]}中恰好有一个会是1，并会对结果产生贡献。因此，梯度计算应该仍然正常进行，结果应该与您上面的代码等价。
