TensorFlow实现zReLU激活函数出现输出形状None报错如何解决

问题描述 #

要实现Nitzan Guberman在2016年发表的论文《On Complex Valued Convolutional Neural Networks》中提出的zReLU激活函数：当输入复数的实部和虚部均为正时输出等于输入，其余场景输出0。

原有实现与报错 #

原有zReLU实现代码：

def zrelu(z: Tensor) -> Tensor:
    angle = tf.math.angle(z)
    return tf.keras.backend.switch(0 <= angle,
                                   tf.keras.backend.switch(angle <= pi / 2,
                                                           z,
                                                           tf.cast(0., dtype=z.dtype)),
                                   tf.cast(0., dtype=z.dtype))

将该激活函数应用到cvnn模型时出现报错：

model = tf.keras.Sequential([
    cvnn.layers.ComplexInput((4)),
    cvnn.layers.ComplexDense(1, activation=tf.keras.layers.Activation(zrelu)),
    cvnn.layers.ComplexDense(1, activation='linear')
])

报错信息为TypeError: unsupported operand type(s) for +: 'NoneType' and 'int'，触发位置为初始化代码行return tf.math.sqrt(6. / (fan_in + fan_out))。

问题原因 #

嵌套tf.keras.backend.switch的两个分支返回值形状不匹配：假值分支返回的是标量0，真值分支返回的是和输入同形状的张量，导致TensorFlow静态形状推导阶段无法确定输出形状，识别为None。即便Activation层自带compute_output_shape方法，也无法覆盖激活函数本身输出形状未知的推导结果，最终下一层参数初始化时读取输入形状为None，触发报错。

解决方案 #

方案1：基于实部虚部判断的掩码实现（最推荐） #

直接对齐原论文的实部、虚部均为正的定义，避免角度计算的精度损失，同时无分支结构对静态图形状推断完全友好：

import tensorflow as tf
from tensorflow import Tensor

def zrelu(z: Tensor) -> Tensor:
    # 直接判断实部、虚部是否均为正，匹配原论文定义
    real_pos = tf.math.real(z) > 0
    imag_pos = tf.math.imag(z) > 0
    cond = tf.logical_and(real_pos, imag_pos)
    # 生成和输入同形状的掩码，相乘后形状完全和输入一致
    mask = tf.cast(cond, dtype=z.dtype)
    return z * mask

方案2：基于角度判断的掩码实现 #

如果需要严格沿用角度判断逻辑，可将分支结构改为掩码乘法，保证输出形状稳定：

import tensorflow as tf
from tensorflow import Tensor
import math

def zrelu(z: Tensor) -> Tensor:
    angle = tf.math.angle(z)
    cond = tf.logical_and(angle >= 0, angle <= math.pi/2)
    mask = tf.cast(cond, dtype=z.dtype)
    return z * mask

方案3：修复原有switch实现 #

如果要保留switch结构，只需将假值分支的标量0改为和输入同形状的零张量即可：

import tensorflow as tf
from tensorflow import Tensor
import math

def zrelu(z: Tensor) -> Tensor:
    angle = tf.math.angle(z)
    zero = tf.zeros_like(z, dtype=z.dtype)
    cond = tf.logical_and(angle >= 0, angle <= math.pi/2)
    return tf.keras.backend.switch(cond, z, zero)

以上三种方案均能让TensorFlow正确识别输出形状，解决报错问题。

内容的提问来源于stack exchange，提问作者J Agustin Barrachina