有什么方法可以加快模型的预测速度吗？

Question

我们目前正在使用 keras 和张量流构建神经网络来评估国际象棋位置。我们遇到了在搜索树中使用的单个样本的预测速度问题。搜索树的用途是检查给定位置上的合法走法，评估每个位置并根据最佳评估得到最佳走法。

对于更高的深度，预测速度会变慢。值得一提的是，我们的神经网络是浅层的 - 3 个 CNN 层和 2 个密集层。我们已经在 CPU 上训练了模型，并且我们也在预测过程中使用 CPU。我们假设，在这个特定的（浅层）网络上它不会影响性能。虽然我们的用例中没有并行性，因此不需要 GPU 计算。

版本：

Python 3.12.4

张量流2.16.1

Keras 3.3.3

我们当时使用 model.predict() 或 model(x) 评估单个样本。我们发现，当用作 Predict_on_batch() 时，其速度与对单个样本进行预测的速度大致相同。我们的目标是尽可能快地获得预测，同时保持当时在单个样本上使用的预测。

我们尝试将模型转换为 TFLite，因为建议获得更好的性能。但由于最新版本不兼容，我们无法转换模型，而降级也不起作用。

我们正在观察每个预测函数在不同批量大小上的速度。

model = keras.models.load_model('firstModel.keras')

print("Durations using model __call__() on small batch")
for i in range(5):
    start = time.time()
    prediction = model(bitboard)
    end = time.time()
    print(end - start)

print("Durations using model.predict_on_batch() on small batch")
for i in range(5):
    start = time.time()
    prediction = model.predict_on_batch(bitboard)
    end = time.time()
    print(end - start)

print("Durations using model.predict() on small batch")
for i in range(5):
    start = time.time()
    prediction = model.predict(bitboard, batch_size=1, verbose=0)
    end = time.time()
    print(end - start)

print("Durations using model.__call__() on larger batch (100 samples)")
for i in range(5):
    start = time.time()
    prediction = model(bitboards)
    end = time.time()
    print(end - start)

print("Durations using model.predict_on_batch() on larger batch (100 samples)")
for i in range(5):
    start = time.time()
    prediction = model.predict_on_batch(bitboards)
    end = time.time()
    print(end - start)

print("Durations using model.predict() on larger batch (100 samples)")
for i in range(5):
    start = time.time()
    prediction = model.predict(bitboards, batch_size=1, verbose=0)
    end = time.time()
    print(end - start)

速度如下：

Durations using model __call__() on small batch
0.055520057678222656
0.007033586502075195
0.006206035614013672
0.007121562957763672
0.005555391311645508
Durations using model.predict_on_batch() on small batch
0.06325101852416992
0.0020132064819335938
0.0010013580322265625
0.0009975433349609375
0.0025305747985839844
Durations using model.predict() on small batch
0.1571955680847168
0.05691671371459961
0.05576348304748535
0.05414080619812012
0.05917525291442871
Durations using model.__call__() on larger batch (100 samples)
0.01164698600769043
0.00638890266418457
0.007528543472290039
0.006807804107666016
0.00751185417175293
Durations using model.predict_on_batch() on larger batch (100 samples)
0.04664158821105957
0.0025255680084228516
0.0010013580322265625
0.0020008087158203125
0.0025064945220947266
Durations using model.predict() on larger batch (100 samples)
0.05106091499328613
0.04923701286315918
0.06421136856079102
0.0651085376739502
0.055069923400878906

什么困扰着我们，但我们不明白如何能够在更短的执行时间内预测更大的批量大小，而不是对单个样本进行预测。我们假设这可能是由于错误的 keras/tensorflow 使用造成的。

主要问题：

有什么建议可以加快预测速度吗？

使用 GPU 运行代码是否有可能提高性能？

您是否会推荐任何其他解决问题的方法或适合我们问题的不同用例？

Answer 1

我将尝试在国际象棋编程的更广泛背景下回答您的问题。

首先，尽管我可能会说显而易见的事情，但 Python 并不是编写国际象棋引擎的最佳编程语言选择。如果您希望提高国际象棋引擎的性能，将其转换为 C++（或任何其他快速语言）绝对值得。如果你不想花很多时间做这件事，也是完全可以理解的。

关于你的问题

我们不明白如何才能在更短的执行时间内对更大的批量大小进行预测，而不是对单个样本进行预测。

这不是因为“错误的 keras/tensorflow 使用”。当您批量运行推理时，您可以避免各种开销，并且操作被矢量化。

但是，假设使用极小极大作为搜索算法，批次不会对你有帮助，因为 alpha beta 剪枝很难并行化，而且它的加速非常重要，你不能没有它。这也代表 GPU 推理（此外，如果您尝试在 GPU 上进行单一推理，数据传输延迟会降低您的性能）。

这就是为什么大多数国际象棋引擎在 CPU 上运行并在单个位置上进行推理的原因。如果你对 GPU 引擎感到好奇，你可以看看 leela chess Zero 是如何工作的。

现在这已经不成问题了，以下是加快神经网络推理速度的方法：

• 通常，onnx 运行时推理比 TensorFlow 或 tflite 推理快得多。您可以期待 5 倍的加速。 为此，您需要将模型转换为

  .onnx

  格式：

    import tf2onnx
    
    # specify the input specs. For example:
    input_spec = (tf.TensorSpec((None, 768), tf.float32, name="input"),) 
    
    output = "onnx_model.onnx"
    
    model_proto, external_tensor_storage = tf2onnx.convert.from_keras(model, opset=15, input_signature=[input_spec], output_path=output)

并运行网络：

from onnxruntime import InferenceSession

data = ... # your input data

sess = InferenceSession("onnx_model.onnx")
input_name = self.sess.get_inputs()[0].name
output_name = self.sess.get_outputs()[0].name
nn_eval = self.sess.run([self.output_name], {self.input_name: data})

• 加速推理的另一种方法是更改模型的拓扑。您可以尝试使用密集网络代替卷积网络，看看是否更快。国际象棋神经网络当前最先进的技术是 NNUE。一开始学习起来很混乱，但速度很快。例如，我观察到，使用 onnx 运行时运行的转换网络的推理时间为 200 微秒，而使用 NNUE 则为 5 微秒。如果您想了解有关 NNUE 的更多信息，我强烈推荐 Stockfish 作者的这篇论文。

如果您想看一下，我在我的国际象棋引擎的自述文件中更详细地讨论了这些速度注意事项。

我希望这有帮助！