视频参与度分类模型Generator训练时损失无法下降排查

我正在做一个参与度检测的视频分类模型，训练到一半就卡壳了——稀疏分类交叉熵损失降到1.06左右就再也降不下去，数据集只有3个类别，这结果跟随机猜没区别。

模型结构 #

模型有两个输入：

• 第一个是10x48x48x1的张量，存的是从10秒视频里采样的10帧48x48灰度图，送进在FER+上预训练的Patt Lite情感识别模型。因为原模型是单图像模型，我加了个聚合层，把10帧挨个喂进去，最后把输出平均一下。
• 第二个是10x329的张量，是用OpenFace每秒提取的面部特征，送进一个前馈神经网络，同样对每秒的输出取平均。
• 最后把两个模型的平均输出拼起来，过一个全连接层得到最终分类结果。

排查情况 #

• 检查了子模型和整体模型的梯度，确认梯度是存在的，没有消失或NaN的情况。
• 做了过拟合测试：模型在手动构造的小数据集上能正常过拟合，但用tf.data.Dataset生成器做的数据集就完全学不进去，训练损失一直停在随机水平。5个epoch后的损失对比明显能看出差异：普通数据集损失掉得很快，生成器数据集基本没变化。

复现代码 #

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Layer, Input, Dense
from keras import models


#Custom attention layer used in the patt lite model. Tensorflow's attention layer was giving us
#nan gradients so we decided to make our own attention layer
class Custom_Attention(tf.keras.layers.Layer):

        def call(self, x):
            size = tf.shape(x)
            x = tf.reshape(x, [size[0], size[1], 1])
            y = tf.einsum("...ij,...jk->...ik", x, tf.transpose(x, perm = [0, 2, 1]))/16.0
            y = tf.matmul(tf.nn.softmax(y), x)
            return tf.reshape(y, [size[0], size[1]])

#We implemented the Patt Lite model described in https://arxiv.org/pdf/2306.09626v1.pdf for our emotion recognition classifier
class Patt_Lite:

    def __init__(self):
        mobile_net = tf.keras.applications.mobilenet.MobileNet(
        include_top = False)


        mobile_net.trainable = False
        inputs = tf.keras.Input(shape = (48, 48, 1), name = "image")
        x = tf.keras.layers.Conv2D(3, (1, 1))(inputs)
        x = tf.keras.layers.Resizing(224, 224)(x)
        x = tf.keras.applications.mobilenet.preprocess_input(x)

        for layer in mobile_net.layers[:56]:
            x = layer(x)

        x = tf.keras.layers.ZeroPadding2D(padding=(1, 1))(x)
        x = tf.keras.layers.Dropout(0.2)(x)
        x = tf.keras.layers.DepthwiseConv2D( (10,10), strides = 2)(x)
        x = tf.keras.layers.Conv2D(256, (1,1), activation = "relu")(x)
        x = tf.keras.layers.Dropout(0.2)(x)
        x = tf.keras.layers.DepthwiseConv2D((3,3))(x)
        x = tf.keras.layers.Conv2D(16, (1,1), activation = "relu")(x)
        x = tf.keras.layers.GlobalAveragePooling2D()(x)
        x = tf.keras.layers.Dense(256, activation = "relu")(x)
        x = Custom_Attention()(x)
        outputs = tf.keras.layers.Dense(9, activation = "softmax")(x)

        self.model = tf.keras.Model(inputs, outputs)
    


#we implement our focus and emotion classifiers as layers in our model
class Emotion_Classifier(Layer):
    def __init__(self, **kwargs):
        super(Emotion_Classifier, self).__init__(**kwargs)

        emoti_model = Patt_Lite().model
        tf.keras.utils.get_custom_objects()["Custom_Attention"] = Custom_Attention
        #we pre-trained the Patt Lite model on the FER+ dataset but for an MRE this isn't necessary
        #emoti_model.load_weights(os.path.join("..", "..", "Models", "Emotion_Rec", "PAtt_Lite_weights.h5"))

        inputs = keras.Input(shape = (48, 48, 1))

        y = emoti_model(inputs)

        #we attach a dense layer to our emotion classifier so that our emotion classifier can be trained
        #to detect engagement
        y = keras.layers.Dense(3, activation = "relu")(y)

        self.model = Model(inputs, y)

    def __call__(self, x):
        return self.model(x)
    


#for our focus classifier, we implemented it as an MLP. Since we're already working with high level features,
#our classifier doesn't need to be as complex as the emotion classifier
class Focus_Classifier(Layer):
    def __init__(self, **kwargs):
        super(Focus_Classifier, self).__init__(**kwargs)
        inputs = keras.Input(shape = (329))
        y = keras.layers.Dense(64, activation = "relu")(inputs)
        y = keras.layers.Dense(64, activation = "relu")(y)
        y = keras.layers.Dense(128, activation = "relu")(y)
        y = keras.layers.Dense(3, activation = "relu")(y)
        self.model = Model(inputs, y)

    def __call__(self, x):
        return self.model(x)
    

#Finally our aggregation layer goes through each of the 10 sampled / averaged frames
#and evaluates the focus and emotion classifiers on each of them
class AggregationLayer(Layer):

    def __init__(self, emo_model, open_model, num_frames, **kwargs):
        super(AggregationLayer, self).__init__(**kwargs)
        self.emoti_model = emo_model
        self.open_model = open_model
        self.num_frames = num_frames

    def __call__(self, inputs):

        #we evaluate our two classifiers on each frame
        emot_outputs = [self.emoti_model(frame) for frame in tf.unstack(inputs[0], axis=1)]
        focus_outputs = [self.open_model(frame) for frame in tf.unstack(inputs[1], axis=1)]

        #average the outputs
        aver_emot_output = tf.reduce_mean(tf.stack(emot_outputs, axis = 1), axis=1)
        aver_focus_output = tf.reduce_mean(tf.stack(focus_outputs, axis = 1), axis=1)

        #and then our output from the aggregation layer are the two 3x1 averaged outputs
        aggregate_output = tf.concat([aver_emot_output, aver_focus_output], axis = 1)
        return aggregate_output

def test_trainable(model, input_size):

    with tf.GradientTape() as tape:
        a = [tf.random.normal(size) for size in input_size]
        z = model(a)
        loss = tf.keras.losses.SparseCategoricalCrossentropy()([1], z)
        #loss = tf.math.reduce_mean()
        #loss = tf.math.reduce_mean(z**2)
    grads = tape.gradient(loss, model.trainable_variables)

    for grad, var in zip(grads, model.trainable_variables):
    #unfortunately all of the gradients are None and I'm not sure why
        try:
            tf.debugging.check_numerics(grad, message=f"{var.name}: ")
            tf.debugging.assert_greater(tf.sum(tf.math.abs(grad)), 0.0)
        except:
            continue


inp_emo = keras.Input((10, 48, 48, 1))
inp_open = keras.Input((10, 329))

emoti_model = Emotion_Classifier()
focus_model = Focus_Classifier()

y = AggregationLayer(emoti_model, focus_model, 10)([inp_emo, inp_open])

outputs = Dense(3, activation = "softmax")(y)

model = Model([inp_emo, inp_open], outputs)

model.summary()


#Test 1: Check if gradients of model exist
test_trainable(Patt_Lite().model,  [(1, 48, 48, 1)])
test_trainable(emoti_model, [(1, 48, 48, 1)])
test_trainable(focus_model, [(1, 329)])
test_trainable(model, [(1, 10, 48, 48, 1), (1, 10, 329)])



#Test 2: Check if model can overfit on small dataset

#I found that the model was unable to overfit on a small dataset made using
#a generator but did overfit on a normal small dataset as can be seen below

#A. Using Generator
def generator():
    feats = (tf.random.normal((1, 10, 48, 48, 1)),
             tf.random.normal((1, 10, 329)))
    labels = tf.reshape(tf.convert_to_tensor( [1], dtype=  tf.int16), [1, 1])

    yield feats, labels

output_signature_train = ( (tf.TensorSpec(shape = (None, 10, 48, 48, 1), dtype = tf.float32),
                      tf.TensorSpec(shape = (None, 10, 329), dtype = tf.float32)),
                    tf.TensorSpec(shape = (None, 1), dtype = tf.int16))

train_ds = tf.data.Dataset.from_generator(generator,
                                          output_signature = output_signature_train)

model.compile(loss = tf.keras.losses.SparseCategoricalCrossentropy(),
              optimizer = keras.optimizers.SGD(learning_rate = 1e-3),
              metrics = ["acc"])

a = train_ds.take(1)

print("Test Using Generator:")
model.fit(a, 
          epochs = 5)

#B. Using Normal Dataset

#we create a new model
model = Model([inp_emo, inp_open], outputs)

model.compile(loss = tf.keras.losses.SparseCategoricalCrossentropy(),
              optimizer = keras.optimizers.AdamW(learning_rate = 1e-2),
              metrics = ["acc"])

#and construct a new dataset without using a generator
feats = (tf.random.normal((1, 10, 48, 48, 1)),
             tf.random.normal((1, 10, 329)))
labels = tf.reshape(tf.convert_to_tensor( [2], dtype=  tf.int16), [1, 1])

print("Test Without Using Generator")
model.fit(feats, labels,
          epochs = 5)

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