tf.function not used for model inference

[mleimeister]

Rasa Open Source version

3.0.x

Rasa SDK version

No response

Rasa X version

No response

Python version

3.8

What operating system are you using?

OSX

What happened?

When running model inference in RasaModel, the predict step being used here is the (eager) Python implementation, not the tf.function created just before that. Therefore, the Tensorflow graph is not used for performance optimization.

Testing a quick fix on the Sara dataset with config from here (limiting training epochs to 20 steps) and timing

rasa test --domain migrated_domain.yml --nlu test --stories test

gave the following inference improvements (tested locally on Macbook Pro 2,4 GHz 8-Core Intel Core i9):

Step	main	tf.function
Core	05:28	01:41
NLU	02:34	00:50
Total	11:18	5:52

Log output:

test-main.log
test-fix.log

Command / Request

#!/bin/zsh
cd /Users/matthias/Workspace/rasa3
git checkout main
cd /Users/matthias/Workspace/training-data/public/Sara

rasa train --domain migrated_domain.yml --data train

{ time rasa test --domain migrated_domain.yml --nlu test --stories test } &> test-main.log

cd /Users/matthias/Workspace/rasa3
git checkout fix-tf-function
cd /Users/matthias/Workspace/training-data/public/Sara

{ time rasa test --domain migrated_domain.yml --nlu test --stories test } &> test-fix.log

Definition of Done

• Implement fix on 3.0.x and 2.8.x branches
• Run regression tests for some datasets/configs to ensure no change in model performance
• Release in upcoming micro versions

[mleimeister]

Results from memory profiling experiments on Carbon bot suggest no overall increase and slightly better utilisation: https://github.com/RasaHQ/research-experiments/tree/kb/tf27/results/2.6vs2.7/results/carbon-bot_tf-function-fix