Keras自定义drawLines层训练报错:transpose期望0维向量输入
首先,你的报错根源在于**tf.py_function破坏了TensorFlow的计算图连续性**:你在自定义drawLines层的call方法中使用tf.py_function调用了依赖numpy、OpenCV的Python函数,这些操作脱离了TensorFlow的自动微分机制,导致优化器(Adadelta)在计算梯度时无法正确处理张量的形状和梯度传递,最终抛出了转置相关的错误。
下面给出两种可行的解决思路,优先推荐第一种(全TensorFlow原生实现):
方案一:用TensorFlow原生操作重写骨架生成逻辑
把所有依赖numpy、OpenCV的代码替换成TensorFlow支持的操作,确保整个计算过程在TensorFlow计算图内完成,这样梯度就能正常回传了。
步骤1:重写关键点坐标提取函数
用TensorFlow的tf.reduce_max和tf.argmax替代numpy的amax和循环查找:
def get_keypointCoordinates(hm): # hm shape: (H, W) max_val = tf.reduce_max(hm) # 找到最大值的位置(先转成一维张量,找索引再转成坐标) flat_idx = tf.argmax(tf.reshape(hm, [-1]), output_type=tf.int32) h_coord = flat_idx // tf.shape(hm)[1] w_coord = flat_idx % tf.shape(hm)[1] # 返回坐标[w, h]和最大值,和原逻辑一致 return tf.stack([w_coord, h_coord]), max_val
步骤2:重写骨架生成函数
用TensorFlow的原生操作实现画直线和圆(替代OpenCV的cv2.line和cv2.circle):
kps_lines = [(0, 1), (1, 2), (2, 6), (7, 12), (12, 11), (11, 10), (5, 4), (4, 3), (3, 6), (7, 13), (13, 14), (14, 15), (6, 7), (7, 8), (8, 9)] def create_skelton(hms): # hms shape: (H, W, num_kps) H, W = tf.shape(hms)[0], tf.shape(hms)[1] num_kps = tf.shape(hms)[2] # 提取所有关键点坐标 kps_coords = [] for kp_idx in tf.range(num_kps): coord, _ = get_keypointCoordinates(hms[:, :, kp_idx]) kps_coords.append(coord) kps_coords = tf.stack(kps_coords) # shape: (num_kps, 2) # 初始化骨架掩码(全白) kp_mask = tf.ones([H, W], dtype=tf.float32) * 255.0 # 绘制每条线段和关键点圆圈 for line in kps_lines: i1, i2 = line p1 = kps_coords[i1] p2 = kps_coords[i2] # 生成线段上的所有点,用tensor_scatter_nd_update绘制直线 num_points = tf.maximum(tf.abs(p1[0]-p2[0]), tf.abs(p1[1]-p2[1])) + 1 t = tf.linspace(0.0, 1.0, num_points) xs = tf.cast(tf.round(p1[0] * (1-t) + p2[0] * t), tf.int32) ys = tf.cast(tf.round(p1[1] * (1-t) + p2[1] * t), tf.int32) indices = tf.stack([ys, xs], axis=1) updates = tf.zeros([num_points], dtype=tf.float32) kp_mask = tf.tensor_scatter_nd_update(kp_mask, indices, updates) # 绘制关键点圆圈(简化为3x3的正方形区域) for p in [p1, p2]: y_min = tf.maximum(p[1]-3, 0) y_max = tf.minimum(p[1]+3, H-1) x_min = tf.maximum(p[0]-3, 0) x_max = tf.minimum(p[0]+3, W-1) ys_circle = tf.range(y_min, y_max+1) xs_circle = tf.range(x_min, x_max+1) y_grid, x_grid = tf.meshgrid(ys_circle, xs_circle, indexing='ij') circle_indices = tf.stack([tf.reshape(y_grid, [-1]), tf.reshape(x_grid, [-1])], axis=1) circle_updates = tf.zeros(tf.shape(circle_indices)[0], dtype=tf.float32) kp_mask = tf.tensor_scatter_nd_update(kp_mask, circle_indices, circle_updates) # 添加通道维度,转为float32 return tf.expand_dims(kp_mask, axis=-1)
步骤3:修改自定义层的call方法
去掉tf.py_function,直接调用TensorFlow原生的函数:
class drawLines(Layer): def __init__(self, output_dim, **kwargs): self.output_dim = output_dim super(drawLines, self).__init__(**kwargs) def build(self, input_shape): super(drawLines, self).build(input_shape) def call(self, inputs): # 对每个batch的样本应用create_skelton return tf.map_fn(create_skelton, inputs, dtype=tf.float32) def compute_output_shape(self, input_shape): output_shape = list(input_shape) output_shape[3] = 1 return tuple(output_shape)
这样修改后,整个自定义层的操作都在TensorFlow计算图内,梯度可以正常回传,训练时就不会出现之前的报错了。
方案二:使用tf.custom_gradient手动定义梯度
如果你确实需要保留numpy和OpenCV的操作(比如某些复杂逻辑难以用TensorFlow实现),可以用tf.custom_gradient装饰image_tensor_func,手动定义梯度的计算方式。不过这种方式需要你自己处理反向传播的梯度,复杂度较高,示例代码如下:
@tf.custom_gradient def image_tensor_func(img4d): # 正向传播逻辑不变,用numpy处理 results = [] img4d_np = img4d.numpy() for img3d in img4d_np: rimg3d = create_skelton(img3d) results.append(rimg3d[:, :, np.newaxis]) results_tensor = tf.convert_to_tensor(results, dtype=tf.float32) # 定义反向传播的梯度函数(这里简单返回全1的梯度,你需要根据实际任务调整) def grad(dy): return dy # 这里只是示例,实际需要根据你的任务定义合理的梯度 return results_tensor, grad
然后自定义层的call方法还是用tf.py_function,但因为有了custom_gradient,梯度可以正常传递。不过注意,这种方式的梯度定义需要符合你的任务逻辑,否则训练效果会受影响。
最后,修改后的模型构建和训练代码可以保持不变,直接运行model.fit即可。
内容的提问来源于stack exchange,提问作者himalayansailor




