ValueError：tf.function 修饰的函数尝试在非首次调用时创建变量

Question

我有一个用模型子类化 API编写的模型tensorflow 2。它包含一个自定义层。问题是，在自定义层中，我需要将输入张量的通道号发送到Conv2D运行时的层。请参阅下面的代码：

自定义层

import tensorflow as tf 

class AuxNet(tf.keras.layers.Layer):
    def __init__(self, ratio=8):
        super(AuxNet, self).__init__()
        self.ratio = ratio
        self.avg = tf.keras.layers.GlobalAveragePooling2D()
        self.max = tf.keras.layers.GlobalMaxPooling2D()

    def call(self, inputs):
        avg = self.avg(inputs)
        max = self.max(inputs)
        avg = tf.keras.layers.Reshape((1, 1, avg.shape[1]))(avg)  
        max = tf.keras.layers.Reshape((1, 1, max.shape[1]))(max)   
        
        # ALERT ---------------------
        input_shape = inputs.get_shape().as_list()
        _, h, w, channels = input_shape
        
        conv1a = tf.keras.layers.Conv2D(channels, kernel_size=1, 
                         strides=1, padding='same',use_bias=True, 
                         activation=tf.nn.relu)(avg)
        
        conv1b = tf.keras.layers.Conv2D(channels, kernel_size=1, strides=1, 
                                    padding='same',use_bias=True, 
                                    activation=tf.nn.relu)(max)
        return tf.nn.sigmoid(conv1a + conv1b)

整个模型

class Net(tf.keras.Model):
    def __init__(self, dim):
        super(Net, self).__init__()
        self.base  = tf.keras.layers.Conv2D(124, 3, 1)
        self.gap   = tf.keras.layers.GlobalAveragePooling2D()
        self.aux   = AuxNet() # init the custom layer
        self.dense  = tf.keras.layers.Dense(128, activation=tf.nn.relu)
        self.out   = tf.keras.layers.Dense(10, activation='softmax')
    
    def call(self, input_tensor, training=False):
        x  = self.base(input_tensor)
        
        # Using custom layer on the input tensor
        aux = self.aux(x)*x

        x = self.gap(aux)
        x = self.dense(x)
        return self.out(x)

正如您所看到的，这些AuxNet类包含Conv2D具有其输入的过滤器大小的层channel。输入只不过是模型类的输入，即Net. 在模型类中初始化自定义图层时，我无法设置其Conv2D图层的通道号。Conv2D因此，我在这里所做的，是在层call的方法中计算通道数AuxNet，我认为这是不好的做法。

这个问题带来了运行时问题。我无法Model以图形模式编译该类，但必须强制启用急切模式。

import numpy as np
import tensorflow as tf 
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Dropout

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()

# train set / data 
x_train = x_train.astype('float32') / 255

# train set / target 
y_train = tf.keras.utils.to_categorical(y_train, num_classes=10)

model = Net((32, 32, 3))

tf.config.run_functions_eagerly(True) # < ----------------

model.compile(
          loss      = tf.keras.losses.CategoricalCrossentropy(),
          metrics   = tf.keras.metrics.CategoricalAccuracy(),
          optimizer = tf.keras.optimizers.Adam())
# fit 
model.fit(x_train, y_train, batch_size=128, epochs=1)

它有效，但训练速度非常慢。但是，如果没有这个，就会出现以下错误，

ValueError：tf.function 修饰的函数尝试在非首次调用时创建变量。

有没有不需要启用急切模式的解决方法？如何有效地将所需的参数传递给这个自定义层？在这种情况下，我不必在call方法中计算通道深度。

Answer 1

我需要了解如何在自定义层中定义内置层。建议所有层都应该初始化该__init__方法。但我们需要channel未知张量的深度，并根据该值filters设置数字。然而，在该build方法中我们可以轻松做到这一点。

class AuxNet(tf.keras.layers.Layer):
    def __init__(self, ratio=8):
        super(AuxNet, self).__init__()
        self.ratio = ratio
        self.avg = tf.keras.layers.GlobalAveragePooling2D()
        self.max = tf.keras.layers.GlobalMaxPooling2D()

    def build(self, input_shape):
        self.conv1 = tf.keras.layers.Conv2D(input_shape[-1], 
                               kernel_size=1, strides=1, padding='same',
                               use_bias=True, activation=tf.nn.relu)
        self.conv2 = tf.keras.layers.Conv2D(input_shape[-1], 
                                               kernel_size=1, strides=1, padding='same',
                                               use_bias=True, activation=tf.nn.relu)
        
        super(AuxNet, self).build(input_shape)
        
    def call(self, inputs):
        avg = self.avg(inputs)
        max = self.max(inputs)
        avg = tf.keras.layers.Reshape((1, 1, avg.shape[1]))(avg)  
        max = tf.keras.layers.Reshape((1, 1, max.shape[1]))(max)   
        
        conv1a = self.conv1(avg)
        conv1b = self.conv2(max)
        
        return tf.nn.sigmoid(conv1a + conv1b)