eval.py

"""
Code adapted from https://github.com/CRIPAC-DIG/GRACE
Linear evaluation on learned node embeddings
"""

import functools

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import f1_score
from sklearn.model_selection import GridSearchCV, train_test_split
from sklearn.multiclass import OneVsRestClassifier
from sklearn.preprocessing import normalize, OneHotEncoder


def repeat(n_times):
    # # repeat 装饰器接受参数 n_times，表示执行被装饰的函数多少次 应用于任何需要进行多次执行并输出统计信息的函数
    def decorator(f):         # decorator 函数接受一个函数 f 作为参数
        @functools.wraps(f)
        def wrapper(*args, **kwargs):                                    # wrapper 函数是实际进行重复执行的函数
            results = [f(*args, **kwargs) for _ in range(n_times)]       # 执行被装饰的函数多次，将结果存储在 results 列表中
            statistics = {}                            # 计算统计信息，包括均值和标准差
            for key in results[0].keys():
                values = [r[key] for r in results]
                statistics[key] = {
                    "mean": np.mean(values),
                    "std": np.std(values),
                }
            print_statistics(statistics, f.__name__)    # 打印统计信息
            return statistics                           # 返回统计信息

        return wrapper

    return decorator


def prob_to_one_hot(y_pred):
    # 将概率预测值转换为独热编码形式的二进制数组
    ret = np.zeros(y_pred.shape, np.bool_)     # 创建一个全零数组，其形状与输入 y_pred 相同，数据类型为布尔类型
    indices = np.argmax(y_pred, axis=1)        # np.argmax 函数获取每行概率预测值中的最大值索引，即最可能的类别
    for i in range(y_pred.shape[0]):           # 通过循环将每行中最可能的类别的位置设置为 True，得到最终的独热编码形式的二进制数组。
        ret[i][indices[i]] = True
    return ret


def print_statistics(statistics, function_name):
    # 打印统计信息的函数
    print(f"(E) | {function_name}:", end=" ")
    for i, key in enumerate(statistics.keys()):        # 遍历统计信息中的每个键值对
        mean = statistics[key]["mean"]
        std = statistics[key]["std"]
        print(f"{key}={mean:.4f}+-{std:.4f}", end="")  # 打印均值和标准差
        if i != len(statistics.keys()) - 1:            # 如果不是最后一个键值对，则打印逗号和空格
            print(",", end=" ")
        else:
            print()


@repeat(3)
# 这个装饰器 @repeat(3) 被应用于名为 label_classification 的函数。这表示 label_classification 函数将会被执行三次，并输出每次执行的统计信息。
def label_classification(embeddings, y, train_mask, test_mask, split="random", ratio=0.1):
    # label_classification 函数，使用 repeat 装饰器进行多次执行
    # 将 PyTorch 张量转换为 NumPy 数组
    X = embeddings.detach().cpu().numpy()
    Y = y.detach().cpu().numpy()
    Y = Y.reshape(-1, 1)

    # 使用 OneHotEncoder 将标签转换为独热编码形式
    onehot_encoder = OneHotEncoder(categories="auto").fit(Y)
    Y = onehot_encoder.transform(Y).toarray().astype(np.bool_)

    # 对特征进行 L2 归一化
    X = normalize(X, norm="l2")

    # 根据 split 参数划分训练集和测试集
    if split == "random":
        X_train, X_test, y_train, y_test = train_test_split(
            X, Y, test_size=1 - ratio
        )
    elif split == "public":
        X_train = X[train_mask]
        X_test = X[test_mask]
        y_train = Y[train_mask]
        y_test = Y[test_mask]

    # 使用 LogisticRegression 进行训练，通过网格搜索选择最佳参数
    logreg = LogisticRegression(solver="liblinear")
    c = 2.0 ** np.arange(-10, 10)

    clf = GridSearchCV(
        estimator=OneVsRestClassifier(logreg),
        param_grid=dict(estimator__C=c),
        n_jobs=8,
        cv=5,
        verbose=0,
    )
    clf.fit(X_train, y_train)

    # 使用训练好的模型进行预测
    y_pred = clf.predict_proba(X_test)
    y_pred = prob_to_one_hot(y_pred)

    # 计算 micro 和 macro F1 分数
    micro = f1_score(y_test, y_pred, average="micro")
    macro = f1_score(y_test, y_pred, average="macro")

    # 返回统计信息，包括 micro 和 macro F1 分数
    return {"F1Mi": micro, "F1Ma": macro}