TensorFlow自定义层输入报错：shape需为int32/int64类型向量

将输入通道、输出通道以tf.constant张量形式传入自定义Keras层类时触发运行错误，报错提示类输入必须是元素类型为{int32,int64}的向量，实际得到shape为[2,1]，错误触发于RandomStandardNormal算子。相关实现代码、输入定义、报错信息如下：

导入依赖库 #

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np

GTLayer类实现代码 #

class GTLayer(keras.layers.Layer):
    
    def __init__(self, in_channels, out_channels, first=True):
        super(GTLayer, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.first = first
        if self.first == True:
            self.conv1 = GTConv(in_channels, out_channels)
            self.conv2 = GTConv(in_channels, out_channels)
        else:
            self.conv1 = GTConv(in_channels, out_channels)
    
    def forward(self, A, H_=None):
        if self.first == True:
            a = self.conv1(A)
            b = self.conv2(A)
            H = tf.matmul(a, b)
            W = [tf.stop_gradient(tf.nn.softmax(self.conv1.weight, axis=1).numpy()),
                 tf.stop_gradient(tf.nn.softmax(self.conv1.weight, axis=1).numpy()) ]
        else:
            a = self.conv1(A)
            H = tf.matmul(H_, a)
            W = [tf.stop_gradient(tf.nn.softmax(self.conv1.weight, axis=1).numpy())]
        return H,W

GTConv层实现代码 #

class GTConv(keras.layers.Layer):
    
    def __init__(self, in_channels, out_channels):
        super(GTConv, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        w_init = tf.random_normal_initializer()
        self.weight = tf.Variable(
            initial_value=w_init(shape=(in_channels, out_channels)),
            trainable=True)
        self.bias = None
        self.scale = tf.Variable([0.1] , trainable=False)
        self.reset_parameters()
        
    def reset_parameters(self):
        n = self.in_channels
        tf.fill(self.weight, 9)
    
    def forward(self, A):
        A = tf.add_n(tf.nn.softmax(self.weight))
        return A 

输入定义 #

inp = tf.constant([4])
out = tf.constant([2])

报错触发代码与错误信息 #

触发代码：

d = GTLayer(inp, out)

错误信息：

--------------------------------------------------------------------------- InvalidArgumentError                      Traceback (most recent call last)  in ()
----> 1 d = GTLayer(inp, out)

5 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name)    7184 def raise_from_not_ok_status(e, name):    7185   e.message += (" name: " + name if name is not None else "")
-> 7186   raise core._status_to_exception(e) from None  # pylint: disable=protected-access    7187     7188

InvalidArgumentError: shape must be a vector of {int32,int64}, got shape [2,1] [Op:RandomStandardNormal]

• 初始化层时传入的in_channels、out_channels是shape为[1]的一维张量，在GTConv中定义权重shape为(in_channels, out_channels)时，两个张量会被组合为shape为[2,1]的二维结构，而RandomStandardNormal算子要求shape参数必须是int32/int64类型标量组成的一维向量，直接传入嵌套张量不符合参数要求，因此触发报错。
• 代码中还存在多处PyTorch迁移TensorFlow的API适配错误，即使解决shape问题也无法正常运行：
  • Keras自定义层的前向传播入口方法名为call，不是PyTorch框架中的forward，方法名写错会导致自定义前向逻辑完全不执行。
  • reset_parameters方法中的tf.fill(self.weight, 9)无法修改变量值：tf.fill是返回新张量的操作，不会原地修改Variable对象，给Variable赋值需要调用assign方法。
  • GTConv前向传播中tf.add_n(tf.nn.softmax(self.weight))写法错误：tf.add_n要求输入是多个张量组成的列表，直接传入单个张量会触发参数错误。
  • GTLayer计算权重列表W时，第二个元素错误引用了conv1的权重，逻辑上应该取conv2的权重。

• 通道参数直接传入Python整数标量，不要传入单元素tf.constant：

inp = 4
out = 2
d = GTLayer(inp, out)

如果业务逻辑要求必须用tf.constant传值，先取出张量中的标量转为Python整数再传入，例如int(inp.numpy())。

• 同步修正其他API适配错误：
  • 将GTLayer、GTConv中的forward方法全部重命名为call
  • 修正reset_parameters的赋值逻辑：

  def reset_parameters(self):
      self.weight.assign(tf.fill(self.weight.shape, 9.0))
  

  • 删除GTConv前向传播中错误的tf.add_n调用，按实际图卷积逻辑调整计算过程
  • 修正GTLayer中W计算的笔误，第二个元素改为tf.stop_gradient(tf.nn.softmax(self.conv2.weight, axis=1).numpy())

内容的提问来源于stack exchange，提问作者Anshuman Sinha