如何为Tensorflow LSTM类编写自定义调用函数?

Question

我已经定义了一个自定义的LSTM层如下:

class LSTMModel(tf.keras.Model):
def __init__(self, CNN_model, num_classes):
super().__init__()
self.cnn_model = CNN_model
self.lstm = tf.keras.layers.LSTM(units=64, return_state=True, dropout=0.3)
self.dense = tf.keras.layers.Dense(num_classes, activation="softmax")
def call(self, input):
pass

然而，我不清楚这里的调用函数中还需要发生什么。我还写了一个通用的CNN类，如下所示:

class generic_vns_function(tf.keras.Model):
# Where would we use layer_units here?
def __init__(self, input_shape, layers, layer_units): 
super().__init__() 
self.convolutions = []
# Dynamically create Convolutional layers and MaxPools
for layer in range(len(layers)):
self.convolutions.append(tf.keras.layers.Conv2D(layer, 3, padding="same", 
input_shape=input_shape, activation="relu"))
# Add MaxPooling layer
self.convolutions.append(tf.keras.layers.MaxPooling2D((2,2)))

# Flatten 
self.flatten = tf.keras.layers.Flatten()

# Dense layer 
self.dense1 = tf.keras.layers.Dense(1024, activation="relu")

def call(self, input):
x = input
for layer in self.convolutions:
x = layer(x)
x = self.flatten(x)
x = self.dense1(x)
return x

但是这里所需的结构对我来说更有意义。我只是初始化了所有的层。我需要做什么来初始化我的LSTM层?

Answer 1

你可以这样写:

import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import LSTM
from tensorflow.keras import Model

class LSTMModel(Model):
def __init__(self, num_classes, num_units=64, drop_prob=0.3):
super().__init__()
self.num_classes = num_classes
self.num_units = num_units
self.drop_prob = drop_prob
self.lstm = LSTM(
units=self.num_units,
return_state=True,
dropout=self.drop_prob)
self.dense = Dense(
num_classes,
activation="softmax")
def call(self, x, training=True):
x, *state = self.lstm(x, training=training)
x = self.dense(x)
return x

然后你可以这样使用:

model = LSTMModel(num_classes=2)
time_series = tf.random.normal((32, 64, 128))
x_pred = model(time_series)
# loss and gradients calculations ...

在初始化自定义层/模型时实例化层，然后在自定义调用实现中通过传递数据来执行它们的call()方法，这是一个常见的tensorflow习惯。