Transform a TF-DF into pure Python code

[rstz]

Sometimes you might want to use a decision tree in an environment that does not allow any new third-party libraries and very minimal code changes. One (hacky, but fun!) solution is to just transform your decision tree into a pure python function - just a sequence of nested if-else statements. Here's some code to do just that:

import tensorflow_decision_forests as tfdf
insp = tfdf.inspector.make_inspector('/path/to/your/model/assets/')
def get_code(node, indent_length=0):
  indent = ' '*indent_length
  if isinstance(node, tfdf.py_tree.node.LeafNode):
    return indent + f"return {node.value.value}\n"
  else:
    code = indent + f"if {node.condition.feature.name} <= {node.condition.threshold}:\n"
    code += get_code(node.pos_child, indent_length+2)
    code += indent + "else:\n"
    code += get_code(node.neg_child, indent_length+2) 
    return code
print(get_code(insp.extract_tree(0).root))

Disclaimers

• This code will only work if all features are numerical. For categorical (or other) features, just adapt the else part in code generator
• Inference in Python using nested if-statements is very slow, but this probably doesn't matter for a single tree
• Individual trees tend to generalize poorly

[rwst]

It sounds like transforming a random forest would be more difficult. How much?

[rstz]

Probably not even that bad. You'd need to create one function per tree (insp.extract_tree(tree_idx) would do the job), and then you write a tiny bit of code to average the tree's predictions. However, random forests can be large and I imagine that at this point the very poor performance would hit you.