使用PyKEEN Pipeline迭代训练时模型分数未提升的问题求助

问题描述 #

使用PyKEEN库实现迭代训练流程，期望每轮基于前一次训练的模型继续训练以提升分数，但首次迭代后分数不再变化，后续迭代中模型参数完全没有更新，训练未生效。

原代码 #

from pykeen.pipeline import pipeline
from pykeen.datasets import Nations
from pykeen.models import TransE

dataset = Nations()
initial_model = TransE
model = initial_model
seed = 0
epochs = 10
budget = 4
scores = []

for i in range(budget):
    test = pipeline(
        dataset=dataset,
        model=model,
        training_kwargs=dict(
            num_epochs=epochs,
            use_tqdm_batch=False,
        ),
        random_seed=seed,
        negative_sampler='basic'
    )

    model = test.model

    scores.append(round(test.metric_results.get_metric("hits_at_10"), 3))

print(scores)

问题原因 #

PyKEEN的pipeline函数默认会对传入的模型进行重新初始化，即使你传入的是已训练的模型实例，它也会重置参数从头训练，导致后续迭代的效果和第一次完全一致，分数没有变化。

解决方法 #

方法1：手动管理训练流程（推荐，更可控） #

放弃使用pipeline，手动拆分数据集加载、模型初始化、训练循环和评估步骤，直接控制模型和优化器的状态：

from pykeen.datasets import Nations
from pykeen.models import TransE
from pykeen.training import SLCWATrainingLoop
from pykeen.evaluation import RankBasedEvaluator
import torch

dataset = Nations()
seed = 0
epochs = 10
budget = 4
scores = []

# 初始化模型、优化器和训练循环
model = TransE(
    triples_factory=dataset.training,
    random_seed=seed,
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
training_loop = SLCWATrainingLoop(
    model=model,
    triples_factory=dataset.training,
    optimizer=optimizer,
    negative_sampler='basic',
    random_seed=seed,
)

evaluator = RankBasedEvaluator()

for i in range(budget):
    # 执行训练
    training_loop.train(
        num_epochs=epochs,
        use_tqdm_batch=False,
    )
    # 评估模型
    results = evaluator.evaluate(
        model=model,
        triples_factory=dataset.testing,
        batch_size=1024,
        use_tqdm=False,
    )
    hits_at_10 = round(results.get_metric("hits_at_10"), 3)
    scores.append(hits_at_10)
    print(f"迭代{i+1}: Hits@10 = {hits_at_10}")

print("最终分数:", scores)

方法2：修改pipeline调用参数 #

如果坚持使用pipeline，需要添加model_kwargs=dict(initialize=False)参数，强制pipeline使用传入的已训练模型的现有参数，不重新初始化：

from pykeen.pipeline import pipeline
from pykeen.datasets import Nations
from pykeen.models import TransE

dataset = Nations()
initial_model = TransE
seed = 0
epochs = 10
budget = 4
scores = []

# 第一次训练：初始化模型
test = pipeline(
    dataset=dataset,
    model=initial_model,
    training_kwargs=dict(
        num_epochs=epochs,
        use_tqdm_batch=False,
    ),
    random_seed=seed,
    negative_sampler='basic'
)
model = test.model
scores.append(round(test.metric_results.get_metric("hits_at_10"), 3))

# 后续迭代：基于现有模型继续训练
for i in range(1, budget):
    test = pipeline(
        dataset=dataset,
        model=model,
        model_kwargs=dict(initialize=False),  # 关键：禁用参数重新初始化
        training_kwargs=dict(
            num_epochs=epochs,
            use_tqdm_batch=False,
        ),
        random_seed=seed,
        negative_sampler='basic'
    )
    model = test.model
    scores.append(round(test.metric_results.get_metric("hits_at_10"), 3))

print(scores)

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