如何将Sklearn随机森林模型转换为原生Python代码？

问题背景 #

我已经训练好Sklearn的随机森林模型，现在需要把它转换成无需调用Sklearn原生predict()方法的原生Python代码，用于嵌入式场景。之前用MicroMLgen生成过C代码，但现在需要Python版本的实现，逻辑类似下面的C示例：

int predict(float *x) {
    uint8_t votes[26] = { 0 };
    // tree #1
    if (x[2] <= 25.0) {
        if (x[1] <= 39.0) {
            if (x[0] <= 268.0) {
                votes[4] += 1;
            } else {
                votes[19] += 1;
            }
        } else {
            if (x[0] <= 302.0) {
                // ... 后续分支逻辑
            }
        }
    }
    // ... 其他树的逻辑
    return 票数最多的类别索引;
}

我的训练代码如下：

import numpy as np
from glob import glob
from os.path import basename
from sklearn.ensemble import RandomForestClassifier
from micromlgen import port

# 加载训练数据集
def load_features(folder):
    dataset = None
    classmap = {}
    for class_idx, filename in enumerate(glob('%s/*.csv' % folder)):
        class_name = basename(filename)[:-4]
        classmap[class_idx] = class_name
        samples = np.loadtxt(filename, dtype=float, delimiter=',')
        labels = np.ones((len(samples), 1)) * class_idx
        samples = np.hstack((samples, labels))
        dataset = samples if dataset is None else np.vstack((dataset, samples))
    return dataset, classmap

if __name__ == '__main__':
    features, classmap = load_features('training_data')
    X, y = features[:, :-1], features[:, -1]
    classifier = RandomForestClassifier(20, max_depth=10).fit(X, y)
    # 生成C++代码
    c_code = port(classifier, classmap=classmap)
    print(c_code)

解决方案 #

通过解析随机森林中每棵树的结构（特征索引、分割阈值、叶节点类别），手动生成对应的Python分支逻辑，模拟投票机制即可实现需求。

1. 解析单棵决策树的函数 #

这个函数递归遍历决策树节点，生成对应分支的Python代码字符串：

def generate_tree_code(tree, feature_names=None, indent=4):
    """
    生成单棵决策树的Python代码
    :param tree: sklearn的决策树对象（classifier.estimators_[i].tree_）
    :param feature_names: 特征名称列表，若为None则用x[0],x[1]...表示
    :param indent: 缩进空格数
    :return: 代码字符串
    """
    indent_str = ' ' * indent
    def recurse(node, depth):
        current_indent = ' ' * (indent * depth)
        # 叶节点：直接投票对应类别
        if tree.children_left[node] == tree.children_right[node]:
            class_idx = np.argmax(tree.value[node])
            return f"{current_indent}votes[{class_idx}] += 1\n"
        # 非叶节点：生成条件分支
        feature_idx = tree.feature[node]
        threshold = tree.threshold[node]
        feature = feature_names[feature_idx] if feature_names else f'x[{feature_idx}]'
        left_code = recurse(tree.children_left[node], depth + 1)
        right_code = recurse(tree.children_right[node], depth + 1)
        code = f"{current_indent}if {feature} <= {threshold}:\n"
        code += left_code
        code += f"{current_indent}else:\n"
        code += right_code
        return code
    return recurse(0, 1)

2. 生成随机森林的完整预测函数 #

遍历森林中所有树，拼接每棵树的代码，生成完整的predict函数：

def generate_random_forest_code(classifier, classmap=None, feature_names=None):
    """
    生成随机森林的原生Python预测代码
    :param classifier: 训练好的RandomForestClassifier对象
    :param classmap: 类别索引到名称的映射，用于生成类别名称转换函数
    :param feature_names: 特征名称列表
    :return: 完整的Python代码字符串
    """
    n_classes = classifier.n_classes_
    # 初始化投票数组
    code = "def predict(x):\n"
    code += f"    votes = [0] * {n_classes}\n"
    # 生成每棵树的分支代码
    for i, tree in enumerate(classifier.estimators_):
        code += f"    # 第{i+1}棵树\n"
        code += generate_tree_code(tree.tree_, feature_names)
    # 计算得票最多的类别
    code += "    # 返回得票最多的类别索引\n"
    code += "    return votes.index(max(votes))\n"
    # 可选：生成返回类别名称的函数
    if classmap:
        code += "\n\ndef predict_class_name(x):\n"
        code += "    class_idx = predict(x)\n"
        code += f"    class_map = {classmap}\n"
        code += "    return class_map[class_idx]\n"
    return code

3. 使用示例 #

在训练代码中添加生成Python代码的逻辑：

if __name__ == '__main__':
    features, classmap = load_features('training_data')
    X, y = features[:, :-1], features[:, -1]
    classifier = RandomForestClassifier(20, max_depth=10).fit(X, y)
    
    # 生成原生Python代码
    python_code = generate_random_forest_code(classifier, classmap)
    # 保存到文件或直接打印
    with open('random_forest_predict.py', 'w') as f:
        f.write(python_code)
    print(python_code)

生成的代码示例 #

最终生成的predict函数大致如下：

def predict(x):
    votes = [0] * 26
    # 第1棵树
    if x[2] <= 25.0:
        if x[1] <= 39.0:
            if x[0] <= 268.0:
                votes[4] += 1
            else:
                votes[19] += 1
        else:
            if x[0] <= 302.0:
                if x[0] <= 115.0:
                    if x[0] <= 86.0:
                        votes[2] += 1
                    else:
                        votes[7] += 1
                # ... 后续分支逻辑
    # 第2棵树
    # ... 其他树的分支逻辑
    # 返回得票最多的类别索引
    return votes.index(max(votes))

说明 #

• 生成的代码完全脱离Sklearn依赖，可直接复制到嵌入式Python环境（如MicroPython）中使用
• 若需要用特征名称代替索引，只需传入feature_names参数即可
• 额外生成的predict_class_name函数可直接返回类别名称，方便业务使用

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