fortest.py

from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import numpy as np

# 构建随机图（示例）
# 图中一共有6个节点，连通的边如下：
# (1, 2): 权值为3
# (1, 3): 权值为4
# (2, 3): 权值为5
# (2, 5): 权值为7
# (3, 4): 权值为2
# (4, 6): 权值为6
graph = {
    (1, 2): 1,
    (1, 3): 2,
    (2, 3): 3,
    (2, 5): 4,
    (3, 4): 5,
    (4, 6): 6
}

# 预处理图数据，将每个边作为一个样本
x = []
y = []
i = 0
for edge, weight in graph.items():
    i += 1
    x.append([edge[0], edge[1], weight])
    if i < 3:
        y.append(1)  # 初始时，每个边都被标记为"选中"
    else:
        y.append(0)

# 划分训练集和测试集
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2)

# 构建随机森林模型，训练并预测
clf = RandomForestClassifier(n_estimators=10)
clf.fit(x_train, y_train)
print(x_test)
y_pred = clf.predict(x_test)

# 计算准确率
accuracy = accuracy_score(y_test, y_pred)
print('准确率：', accuracy)

# 获取预测结果（即每条边的选取情况）
result = clf.predict(x)
result = np.array(result, dtype=bool)

# 输出选取结果
for i, edge in enumerate(graph.keys()):
    if result[i]:
        print('选中：', edge)
    else:
        print('未选中：', edge)

    import pickle

    # 假设模型对象为clf
    with open('model.pkl', 'wb') as f:
        pickle.dump(clf, f)

    print("保存完成")

    # 从文件model.pkl中加载模型对象
    with open('model.pkl', 'rb') as f:
        clf = pickle.load(f)

    # 使用加载回来的模型，进行预测等操作
    y_pred = clf.predict(x_test)