众所周知，一般 research 里面所使用的炼丹框架以 PyTorch 居多。然而在工业界的模型部署中，TensorFlow 因能较好地支持多 GPU 的特性而广受青睐。

其中，Keras 以 TensorFlow 为后端（之前的 Keras 版本也能以 Theano、MXNet 等为后端），对 TensorFlow 的诸多 API 进行了封装，为模型的部署进一步降低了门槛。并且，TensorFlow 也把 Keras 集成进来，作为 TensorFlow 的一部分。

俺也是最近才系统性地接触到 Keras，之前都是 PyTorch 居多，用 TensorFlow 或者 Keras 仅仅是在复现其他人的工作上会用到。所以这不可避免地就会踩到很多坑。在这里记录下来，以方便之后遇到类似坑的时候，能快速反应下来是出现了什么问题。本人机器配置如下：

  • 操作系统：Windows 11 家庭中文版 21H2
  • 显卡：NVIDIA GeForce RTX 3060
  • Python 版本：3.9.7
  • CUDA 版本：11.6（可能就是因为这个 11.6 带来了之后的一系列巨坑……现在跑去 NVIDIA 官网下 CUDA，它会首退给你 11.6，但是未必最新的就是最好的，不想折腾的话还是用 11.2 之类的稳定版本……）
  • TensorFlow 版本：2.8.0
  • Keras 版本：2.8.0

keras.Sequential 中 model.build 的问题

报错信息如下：

1
2
return int(fan_in), int(fan_out)
TypeError: int() argument must be a string, a bytes-like object or a number, not 'NoneType'

我们以一个简单的 LSTM 为例，用keras.Sequential来构造：

1
2
3
4
5
6
model = keras.Sequential()
model.add(keras.layers.LSTM(64, return_sequences=False))
model.add(keras.layers.Dense(32, activation="ReLU"))
model.add(keras.layers.Dense(1))
model.build(input_shape=(None, 10, 12))
print(model.summary())

这里我们也就成功通过build，构建了一个序列长度为 10，特征维度为 12 的 LSTM。并且可以在后面进行正常的训练和测试。model.summary()结果如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
Model: "sequential"
_________________________________________________________________
 Layer (type)                Output Shape              Param #
=================================================================
 lstm (LSTM)                 (None, 64)                19712

 dense (Dense)               (None, 32)                2080

 dense_1 (Dense)             (None, 1)                 33

=================================================================
Total params: 21,825
Trainable params: 21,825
Non-trainable params: 0
_________________________________________________________________

但是，在模型的保存和读取测试中，出现了上述的报错。

1
2
3
model.save('./checkpoint/test.h5')
time.sleep(1)
model = keras.models.load_model('./checkpoint/test.h5')

经过查询，用 Sequential 的话，我们在 LSTM 的定义中要确定好input_shape，而非在最后用model.build()。具体细节可以看到下面的 Stack Overflow：

https://stackoverflow.com/questions/64829642/keras-load-model-typeerror-int-argument-nonetype

里面提到了可以用model.to_json()查看模型的细节。具体操作的结果是：

  • model.build()之后，input_shape为[null, 10, 12]

但是在模型完成了训练和验证之后，也就是在model.save()之前，input_shape变成了[null, null, null]。这就确实是 Stack Overflow 的问题。直接在 LSTM 的定义里面说明input_shape=(10, 12)就能解决问题。

1
2
3
4
5
model = keras.Sequential()
model.add(keras.layers.LSTM(64, return_sequences=False, input_shape=(10, 12)))
model.add(keras.layers.Dense(32, activation="ReLU"))
model.add(keras.layers.Dense(1))
print(model.summary())

训练时 checkpoint 回调不停地报关于 LSTMcell 的 warning

这个应该是没有将存储路径写成 hdf5 的后缀。直接改一下就行了

1
2
3
4
checkpoint_path = './checkpoint/' # NO
checkpoint_path = './checkpoint/weights.{epoch:02d}.h5' # YES
cp_callback = keras.callbacks.ModelCheckpoint(
            checkpoint_path, save_weights_only=False, save_best_only=True, monitor='val_loss')

当 model.fit 的训练集使用 generator 时停不下来

众所周知，model.fit方法能接受的数据集可以是list、np.array、pd.DataFrame，其中后两者都会默认第一维为 batch 的索引。

然而，在实际的应用场景中，我们的数据集本身可能比较大，没法一次性装到内存中，于是model.fit也支持我们用generator作为模型的输入。如下就是一个简单的generator例子：

1
2
3
4
5
6
7
8
9
10
11
def train_generator():
    while True:
        batch = []
        for i in range(10):
            batch.append(i)
            if (len(batch) == 3):
                yield batch
                batch = []
        if (len(batch) > 0):
            yield batch
            batch = []

这就会得到一个generator，并且可以通过__next__()方法得到一个批次的训练数据。

然后送到model.fit中：

1
2
3
model.fit(train_generator,
          epochs=10,
          validation_data=val_generator)

这样会发现模型会一直在第一个 epoch 中训练，并且预估还需时间为nan。这是因为，如此定义的 generator 会一直产生样本。model.fit 在外部调用时并不清楚该 generator 会生成多少样本。

所以，我们还需要在model.fit中加入一个steps_per_epoch参数，用来指示一轮有多少个批次。这可以通过简单的计算得来（要保留最后一个批次就取上整，不保留最后一个批次就取下整，参考DataLoader中的drop_last参数）。类似地，加入一个validation_steps来指示验证集一轮有多少个批次。

1
2
3
4
5
model.fit(train_generator,
          epochs=20,
          steps_per_epoch=steps_per_epoch,
          validation_data=val_generator,
          validation_steps=validation_steps)

当 model.fit 的训练集使用 generator 时不能多线程取数据

当我们需要用到generator来作为训练集和验证集的输入时，一般都是因为不能直接将一整个训练集装进内存（否则，我们完全可以用更简单的np.array来做）。此时，我们 CPU 的计算资源往往不会是计算瓶颈。因此，我们寄希望于利用多线程读取训练数据的方式，来加速训练。

1
2
3
4
5
6
model.fit(train_generator,
          epochs=20,
          steps_per_epoch=steps_per_epoch,
          validation_data=val_generator,
          validation_steps=validation_steps,
          workers=4)

然而，这会导致报错，因为 Keras 无法信任我们传进来的 generator 是线程安全的。所以，我们需要另谋他法。

这里我们用到的是继承keras.utils.Sequence这个类。和 PyTorch 里面我们继承torch.utils.data.Dataset一样，我们也需要重写一些方法：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class DataGenerator(keras.utils.Sequence):
    def __init__(self, data, batch_size) -> None:
        super(DataGenerator, self).__init__()

        self.data = data
        self.batch_size = batch_size
        self.batch_num = (data.shape[0] + batch_size - 1) // batch_size
        self.on_epoch_end()

    def __len__(self):
        return self.batch_num

    def __getitem__(self, index):
        data_batch = self.data[index*self.batch_size:(index+1)*self.batch_size]
        X_batch = data_batch[..., :-1]
        y_batch = data_batch[..., -1]

        # 返回一个批次
        return X_batch, y_batch

    def on_epoch_end(self):
        # 训练集和测试集在每轮训练（测试）的随机性来源于此方法的实现
        np.random.shuffle(self.train_dataset)

于是我们就可以愉快地实现训练时利用多线程取数据了，并且也不需要steps_per_epoch和validation_steps参数（一轮多少批次已经由__len__()方法得到）：

1
2
3
4
model.fit(train_generator,
          epochs=20,
          validation_data=val_generator,
          workers=4)

cuDNN 不能加载 dll

在使用 LSTM 的时候，一般不会出现这个问题。

但是在使用卷积层的时候，例如Conv1D，可能会遇到如下报错：

1
Could not load library cudnn_cnn_infer64_8.dll. Error code 193

但是我已经明明将 cuDNN 中的所有文件复制到 CUDA 文件夹下了，仍然找不到。

后面折腾了一通，发现 NVIDIA 官网推 cuDNN 也是推的最新版 8.4.0。但是似乎 TensorFlow 和 Keras 并不支持这么新的 cuDNN。所以才会报错。

当我们再次查询 TensorFlow 的推荐 CUDA 和 cuDNN 版本：

https://www.tensorflow.org/install/source_windows?hl=en#gpu

发现推荐的 CUDA 和 cuDNN 版本分别为 11.2 和 8.1。

所以我们重新到 NVIDIA 官网上下一个低版本的 cuDNN。看到 8.2.0 是最低的支持 11.x 的 cuDNN 版本，果断下载、覆盖。嘿，您猜怎么着？好了！

所以关键就是 CUDA 和 cuDNN 版本不是越新越好，还是得考虑与 TensorFlow 和 Keras 的适配性。

当然 30 系显卡貌似只能装 11.x 的 CUDA，是时候和 TensorFlow 1.x say goodbye 了。

和所有的 1.x 说拜拜
和所有的 2.x 说嗨嗨

关于 XLA

在模型的编译中，我们可以在官方文档看到有个 jit_compile 的可选项，其取值默认为 auto。

看了下面的描述，我一开始以为 jit_compile=True 能提高 GRU 模型的效率，但是后面发现会有很多问题：

  1. 模型的运算效率上并没有得到优化，反而还在做负优化；

  2. 模型可能会意外触发报错。这是我在数据迭代的时候打印出当前数据的 index，结果如下：

1
2
3
4
5
6
7
index = 45278
index = 45279
2024-04-03 16:50:32.020688: I tensorflow/stream_executor/cuda/cuda_blas.cc:1614] TensorFloat-32 will be used for the matrix multiplicat.;ion. This will only be logged once.
index = 45280
corrupted double-linked list
index = 45281
index = 45282

然后训练报错推出。当时想着纯 python 代码，哪里有可能涉及到啥双向链表，那必然是 keras 或者 tensorflow 的底层代码弄出来的活儿。那哥们这半吊子也不可能真去嗯啃 tensorflow 的 C++ 源码吧。关键是还在那调半天还没有发现问题，后面才问了工友，结果工友发现了这问题。

另一次报错的内容更加直接：

1
2
3
4
5
2024-04-07 18:04:00.859694: E tensorflow/compiler/xla/status_macros.cc:57] INTERNAL: RET_CHECK failure (tensorflow/compiler/xla/service/gpu/gpu_executable_run_options.cc:72)
 0 <= local_device_ordinal && local_device_ordinal < gpu_global_device_ids->size()
*** Begin stack trace ***
        tensorflow::CurrentStackTrace[abi:cxx11]()
        xla::status_macros::MakeErrorStream::Impl::GetStatus()
  1. 在训练的过程中，模型可能发生内存泄漏。具体的发现是通过加了一个回调来每轮开始和结束的内存使用量：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# https://github.com/keras-team/keras/issues/15887#issuecomment-1010096128
import os
import psutil
from tensorflow import keras
from loguru import logger


class MemoryUsageCallback(keras.callbacks.Callback):
    '''
    Monitor memory usage on epoch begin and end.
    '''

    def on_epoch_begin(self, epoch, logs=None):
        logger.debug('**Epoch {}**'.format(epoch+1))
        logger.debug('Memory usage on epoch begin: {}'.format(psutil.Process(os.getpid()).memory_info().rss))  # nopep8

    def on_epoch_end(self, epoch, logs=None):
        logger.debug('Memory usage on epoch end:   {}'.format(psutil.Process(os.getpid()).memory_info().rss))  # nopep8

然后在训练的时候带上这个 callback 即可：

1
2
3
4
5
6
7
memory_usage_callback = MemoryUsageCallback()
callbacks = [
    early_stopping_callback,
    reduce_lr_on_plateau_callback,
    model_checkpoint_callback,
    memory_usage_callback
]

就可以监控到模型训练过程中每个 epoch 开始和结束时的内存使用情况。

不多说了，看看下面的训练日志吧：

可以看到内存使用量几乎随 epoch 线性增长！这是多个 epoch 训练时绝对不能接受的！（事实上我们期待内存使用量应该是恒定的）

然后我去找了下相关 keras 的 issue：

https://github.com/keras-team/keras/issues/19071

_这里多一句嘴，把 XLA 关了之后，内存占用也会随 epoch 增加有些微的增长，虽然无伤大雅，但是也表明说不定哪里就藏着一些内存泄露……_

后面工友对 XLA 的经验是：

之前测过这玩意，基本上没啥用，没想到还导致性能倒退[doge]

所以他对我所给出的人生经验就是：

对于不熟悉的功能/框架/特性，需要谨慎引入

更改模型里面层的名称

这个事情本来的需求是想在模型构建的时候嵌入一个 metadata，但是呢 keras 模型本身是没有预留一个类似于 metadata 之类的字段的。后面发现如果是只要嵌一个 metadata 的话，其实就直接嵌入到层的 name 字段里面就行了。于是某些 stackoverflow 上的解决方法就是：

之后如果需要增加 metadata 内容的话还得通过另外 dump 个 pickle 对象的方式来搞