TCN训练动作分割模型时损失为NaN问题排查

任务基本信息 #

• 运行环境：Google Colab平台
• 任务目标：视频动作分割，共包含51类动作
• 模型结构：时间卷积网络（Temporal Convolutional Network, TCN）
• 输入数据：预提取的视频帧特征，单帧特征维度为1024，单个视频的特征形状为(x, 1024)，x为对应视频的总帧长，不同视频的帧长存在差异
• 数据集拆分：按7:3比例划分为训练集（占70%）、验证集（占30%）；每轮训练使用全部训练样本，训练时打乱批次顺序，避免模型学习到固定的样本排列规律
• 数据加载方案：继承Keras Sequence类编写自定义DataGenerator，保证单个训练批次仅对应单视频内的一段连续序列，不同视频的序列无交叉重叠

现有实现代码 #

!pip install -q keras-tcn --no-dependencies
from tcn import TCN, tcn_full_summary

class DataGenerator(Sequence):

  def __init__(self, list_IDs, labels, batch_size=None, dim=None, n_channels=None,
              n_classes=uniques.shape[0], seq_length=max_length, shuffle=True):
      'Initialization'
      self.dim = dim
      self.batch_size = batch_size
      self.labels = labels
      self.list_IDs = list_IDs
      self.n_channels = n_channels
      self.n_classes = n_classes
      self.shuffle = shuffle
      self.seq_length = seq_length
      self.on_epoch_end()

  def on_epoch_end(self):
    'Updates indexes after each epoch'
    self.indexes = np.arange(len(self.list_IDs))
    if self.shuffle == True:
        np.random.shuffle(self.indexes)
    
  def __data_generation(self, video_index, first_frame):
    'Generates data containing batch_size samples' # X : (n_samples, *dim, n_channels)
    # Initialization
    X = np.empty((self.batch_size, *self.dim))
    y = np.empty((self.batch_size), dtype=int)

    ID = self.list_IDs[video_index]
    
    data = np.load('/content/drive/MyDrive/features/trimmed_rgb-' + ID + '_rgb.npy')
    num_frames = data.shape[0]

    for frame in range(first_frame, first_frame + self.batch_size):
      if frame >= num_frames:
        break
      
      X[frame - first_frame] = data[frame]
      y[frame - first_frame] = self.labels[ID][frame]
        
    return X.reshape(1, * X.shape), tf.reshape(to_categorical(y, num_classes=self.n_classes), [1, self.batch_size, self.n_classes])

  @property
  def num_videos(self):
    return len(self.list_IDs)

  def __len__(self):
    'Denotes the number of batches per epoch'
    return int(np.floor(self.seq_length / self.batch_size))

  def __getitem__(self, index):
    'Generate one batch of data'
    # Generate indexes of the batch
    video_index = int((index * self.batch_size) / self.seq_length)
    frame_index = int((index * self.batch_size) % self.seq_length)
    
    # Generate data
    X, y = self.__data_generation(video_index, frame_index)

    return X, y

params = {'dim': [1024],
      'batch_size': 32,
      'n_classes': uniques.shape[0],
      'n_channels': None,
      'seq_length': sum_train_lengths,
      'shuffle': True}

# Generators
training_generator = DataGenerator(partition['train'], labels, **params)

params = {'dim': [1024],
          'batch_size': 32,
          'n_classes': uniques.shape[0],
          'n_channels': None,
          'seq_length': sum_valid_lengths,
          'shuffle': True}

validation_generator = DataGenerator(partition['validation'], labels, **params)

model = keras.models.Sequential()
model.add(TCN(input_shape=(None, 1024), nb_filters=64, kernel_size=3, nb_stacks=2,return_sequences=True))
model.add(tf.keras.layers.Dense(uniques.shape[0],activation='softmax'))

# Compile the model
model.compile(loss='categorical_crossentropy', optimizer='adam',metrics="accuracy")

# Train model on dataset
history = model.fit(training_generator,validation_data=validation_generator,epochs=30)

异常现象 #

1. 初始版本未对标签y执行tf.reshape操作时，运行触发形状不兼容错误：

ValueError: Shapes (None, None) and (None, None, 51) are incompatible

2. 添加tf.reshape操作后模型可正常启动训练，但训练集与验证集的损失值均为NaN；已排查training_generator.__len__()覆盖范围内的训练数据，确认原始特征与标签数据本身不存在NaN值。

咨询问题 #

• 排除原始数据存在NaN的情况后，还有哪些原因会导致训练过程中损失变为NaN？
• 当前代码中使用的tf.reshape操作是否会引发NaN问题？

初步排查结论 #

初步判断问题根源在于__data_generation方法中的break逻辑会生成缺失值：初始化X和y数组时使用np.empty申请内存，未填充初始值，当触发break提前终止帧遍历逻辑时，数组末尾未被赋值的位置会保留内存中的随机垃圾值，导致生成器__len__覆盖的批次范围内存在部分填充不完整的样本；同时该逻辑会直接丢弃每个视频长度不足一个batch的末尾序列，并非最优的训练实现方式。待重构数据生成与拟合逻辑解决问题后，会进一步更新解决方案。

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