使用TensorFlow高层API自定义Layer时出现'_expects_mask_arg'属性错误

嘿，我来帮你搞定这个问题！首先，那个AttributeError的直接原因很容易修复：你在build方法里写错了父类的类名。

看这段代码：

def build(self,input_shape):
    super(MyLayer, self).build(input_shape)

这里的MyLayer应该是你的自定义层类名Reconstruction_Layer才行。Keras的父类build方法会初始化一些内部必要属性（包括_expects_mask_arg），如果调用了错误的父类，这些属性就不会被正确设置，自然就抛出了这个错误。

不过修复这个之后，你的代码还有几个在TensorFlow 2.x中需要优化的问题，我一起帮你梳理和修正，让它更符合TF的最佳实践：

主要问题和修复点 #

• 错误的父类super调用：刚才提到的，把MyLayer改成Reconstruction_Layer
• 用numpy数组存储中间状态：在Keras层里，numpy数组没法参与自动微分，也不兼容图模式，应该全部改用张量操作
• 嵌套循环效率太低：手动写三层for循环在TF里会跑得巨慢，必须用向量化操作替代
• 不必要的类变量初始化：像self.reconstructed_img、self.normal_light这些不需要作为类成员，在call方法里用张量创建就行

修正后的完整代码 #

import math
import tensorflow as tf

class Reconstruction_Layer(tf.keras.layers.Layer):
    def __init__(self):
        super(Reconstruction_Layer, self).__init__()
        # 删掉了所有提前初始化的numpy数组，改用张量操作
    
    def build(self, input_shape):
        # 修正super的类名，确保父类build方法正确执行
        super(Reconstruction_Layer, self).build(input_shape)
    
    def call(self, input_layer):
        normal, albedo, light = input_layer
        
        # 向量化计算球谐函数，彻底替换嵌套循环
        x = normal[..., 0]
        y = normal[..., 1]
        z = normal[..., 2]
        
        # 计算各个球谐分量，保持张量形状匹配
        sphar0 = tf.fill(tf.shape(normal)[0:3], 1 / math.sqrt(4 * math.pi))
        sphar1 = math.sqrt(3/(4*math.pi)) * z
        sphar3 = math.sqrt(3/(4*math.pi)) * y
        sphar4 = 0.5 * math.sqrt(5/(4*math.pi)) * (3 * tf.square(z) - 1)
        sphar5 = 3 * math.sqrt(5/(12*math.pi)) * z * x
        sphar6 = 3 * math.sqrt(5/(12*math.pi)) * z * y
        sphar7 = 1.5 * math.sqrt(5/(12*math.pi)) * (tf.square(x) - tf.square(y))
        sphar8 = 3 * math.sqrt(5/(12*math.pi)) * x * y
        
        # 拼接所有球谐分量，得到形状为[None, 128, 128, 9]的张量
        # 注意你原来的代码里sphar[2]没赋值，这里用tf.zeros_like补全
        normal_light = tf.stack([sphar0, sphar1, tf.zeros_like(sphar0), sphar3, sphar4, sphar5, sphar6, sphar7, sphar8], axis=-1)
        
        # 调整light的形状，方便后续点积计算：从[27]变成[1,1,1,3,9]
        light_reshaped = tf.reshape(light, (1,1,1,3,9))
        # 计算每个颜色通道的点积，利用广播机制自动匹配维度
        dot_product = tf.reduce_sum(normal_light[..., tf.newaxis, :] * light_reshaped, axis=-1)
        # 乘以albedo得到最终的重建图像
        reconstructed_img = albedo * dot_product
        
        return reconstructed_img

简单测试示例 #

你可以这样测试修复后的层：

# 生成示例输入张量
normal_input = tf.random.normal((1, 128, 128, 3))
albedo_input = tf.random.normal((1, 128, 128, 3))
light_input = tf.random.normal((27,))

# 实例化层并调用
recon_layer = Reconstruction_Layer()
output = recon_layer([normal_input, albedo_input, light_input])
print(output.shape)  # 应该输出 (1, 128, 128, 3)，符合预期

额外说明 #

• 我注意到你原来的代码里sphar[2]没有赋值，所以在修正后的代码里用tf.zeros_like(sphar0)补全了这个分量，如果这不是你想要的，可以根据实际需求修改
• 向量化操作不仅提升了运行速度，还能确保代码兼容TensorFlow的图模式和自动微分功能，这对训练模型来说是必需的

内容的提问来源于stack exchange，提问作者Trilok Padhi