我使用tensorflow 2
中的模型子类化API编写了一个模型。该模型包含一个自定义层。问题在于,在自定义层中,我需要在运行时将输入张量的通道数传递给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) # 警告 --------------------- 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() # 初始化自定义层 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) # 在输入张量上使用自定义层 aux = self.aux(x)*x x = self.gap(aux) x = self.dense(x) return self.out(x)
如您所见,AuxNet
类包含Conv2D
层,其过滤器大小为其输入的channel
。而输入正是模型类Net
的输入。在模型类中初始化自定义层时,我无法设置其Conv2D
层的通道数。因此,我在AuxNet
层的call
方法中计算这个Conv2D
的通道数,我认为这是不好的做法。
这个问题带来了运行时问题。我无法在图模式下编译Model
类,而被迫启用急切模式(eager mode)。
import numpy as npimport tensorflow as tf from tensorflow.keras.models import Modelfrom tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Dropout(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()# 训练集/数据 x_train = x_train.astype('float32') / 255# 训练集/目标 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(x_train, y_train, batch_size=128, epochs=1)
它可以工作,但训练速度非常慢。然而,如果不这样做,会出现以下错误,
ValueError: tf.function-decorated function tried to create variables on non-first call.
有没有不需要启用急切模式的解决方法?我如何有效地将所需参数传递给这个自定义层?这样的话,我就不必在call
方法中计算通道深度了。
回答:
我需要查看如何在自定义层内定义内置层。建议所有层都应该在__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)