Trainer¶

class unike.config.Trainer(model: Strategy | None = None, data_loader: torch.utils.data.DataLoader | None = None, epochs: int = 1000, lr: float = 0.5, opt_method: str = 'Adam', use_accelerator: bool = False, use_gpu: bool = True, device: str = 'cuda:0', tester: Tester | None = None, test: bool = False, valid_interval: int | None = None, log_interval: int | None = None, save_interval: int | None = None, save_path: str | None = None, use_early_stopping: bool = True, metric: str = 'hits@10', patience: int = 2, delta: float = 0, wandb_logger: WandbLogger | None = None)[源代码]¶

主要用于 KGE 模型的训练。

例子:

from unike.data import KGEDataLoader, BernSampler, TradTestSampler
from unike.module.model import TransE
from unike.module.loss import MarginLoss
from unike.module.strategy import NegativeSampling
from unike.config import Trainer, Tester

# dataloader for training
dataloader = KGEDataLoader(
        in_path = "../../benchmarks/FB15K/", 
        batch_size = 8192,
        neg_ent = 25,
        test = True,
        test_batch_size = 256,
        num_workers = 16,
        train_sampler = BernSampler,
        test_sampler = TradTestSampler
)

# define the model
transe = TransE(
        ent_tol = dataloader.get_ent_tol(),
        rel_tol = dataloader.get_rel_tol(),
        dim = 50, 
        p_norm = 1, 
        norm_flag = True)

# define the loss function
model = NegativeSampling(
        model = transe, 
        loss = MarginLoss(margin = 1.0),
        regul_rate = 0.01
)

# test the model
tester = Tester(model = transe, data_loader = dataloader, use_gpu = True, device = 'cuda:1')

# train the model
trainer = Trainer(model = model, data_loader = dataloader.train_dataloader(),
        epochs = 1000, lr = 0.01, use_gpu = True, device = 'cuda:1',
        tester = tester, test = True, valid_interval = 100,
        log_interval = 100, save_interval = 100,
        save_path = '../../checkpoint/transe.pth', delta = 0.01)
trainer.run()

__init__(model: Strategy | None = None, data_loader: torch.utils.data.DataLoader | None = None, epochs: int = 1000, lr: float = 0.5, opt_method: str = 'Adam', use_accelerator: bool = False, use_gpu: bool = True, device: str = 'cuda:0', tester: Tester | None = None, test: bool = False, valid_interval: int | None = None, log_interval: int | None = None, save_interval: int | None = None, save_path: str | None = None, use_early_stopping: bool = True, metric: str = 'hits@10', patience: int = 2, delta: float = 0, wandb_logger: WandbLogger | None = None)[源代码]¶

创建 Trainer 对象。

参数:

model (unike.module.strategy.Strategy) – 包装 KGE 模型的训练策略类
data_loader (torch.utils.data.DataLoader) – torch.utils.data.DataLoader
epochs (int) – 训练轮次数
lr (float) – 学习率
opt_method (str) – 优化器: ‘Adam’ or ‘adam’, ‘Adagrad’ or ‘adagrad’, ‘SGD’ or ‘sgd’
use_accelerator (bool) – 使用 accelerate 进行分布式训练
use_gpu (bool) – 是否使用 gpu
device (str) – 使用哪个 gpu
tester (unike.config.Tester) – 用于模型评估的验证模型类
test (bool) – 是否在测试集上评估模型, tester 不为空
valid_interval (int) – 训练几轮在验证集上评估一次模型, tester 不为空
log_interval (int) – 训练几轮输出一次日志
save_interval (int) – 训练几轮保存一次模型
save_path (str) – 模型保存的路径
use_early_stopping (bool) – 是否启用早停，需要 tester 和 save_path 不为空
metric (str) – 早停使用的验证指标，可选值：’mr’, ‘mrr’, ‘hits@N’, ‘mr_type’, ‘mrr_type’, ‘hits@N_type’。默认值：’hits@10’
patience (int) – unike.utils.EarlyStopping.patience 参数，上次验证得分改善后等待多长时间。默认值：2
delta (float) – unike.utils.EarlyStopping.delta 参数，监测数量的最小变化才符合改进条件。默认值：0
wandb_logger (unike.utils.WandbLogger) – unike.utils.WandbLogger 对象

__weakref__¶: list of weak references to the object (if defined)

accelerator: Accelerator¶: accelerate.Accelerator 对象

configure_optimizers()[源代码]¶: 可以通过重新实现该方法自定义配置优化器。

data_loader: torch.utils.data.DataLoader¶: __init__() 传入的 torch.utils.data.DataLoader

delta: float¶: unike.utils.EarlyStopping.delta 参数，监测数量的最小变化才符合改进条件。默认值：0

device: Union[torch.device, str]¶: gpu，利用 device 构造的 torch.device 对象

early_stopping: EarlyStopping¶: 早停对象

epochs: int¶: epochs

get_device() → torch.device | str[源代码]¶

返回当前进程的设备。

返回:: 设备信息
返回类型:: Union[torch.device, str]

get_model() → Model[源代码]¶

返回原始的 KGE 模型。

返回:: KGE 模型
返回类型:: unike.module.model.Model

is_local_main_process() → bool[源代码]¶

当前进程是否是主进程。

返回:: 当前进程是否是主进程。
返回类型:: bool

log_interval: int | None¶: 训练几轮输出一次日志

lr: float¶: 学习率

metric: str¶: 早停使用的验证指标，可选值：’mr’, ‘mrr’, ‘hits@N’, ‘mr_type’, ‘mrr_type’, ‘hits@N_type’。默认值：’hits@10’

model: Strategy¶: 包装 KGE 模型的训练策略类，即 unike.module.strategy.Strategy

opt_method: str¶: 用户传入的优化器名字字符串

optimizer: torch.optim.SGD | torch.optim.Adagrad | torch.optim.Adam | None¶: 根据 __init__() 的 opt_method 生成对应的优化器

patience: int¶: unike.utils.EarlyStopping.patience 参数，上次验证得分改善后等待多长时间。默认值：2

print_test(sampling_mode: str, epoch: int = 0)[源代码]¶

根据 tester 类型进行链接预测。

参数:: sampling_mode (str) – 数据

run()[源代码]¶: 训练循环，首先根据 use_gpu 设置 model 是否使用 gpu 训练，然后根据 opt_method 设置 optimizer，最后迭代 data_loader 获取数据，并利用 train_one_step() 训练。

save_interval: int | None¶: 训练几轮保存一次模型

save_path: str | None¶: 模型保存的路径

scheduler: torch.optim.lr_scheduler.MultiStepLR | None¶: 学习率调度器

test: bool¶: 是否在测试集上评估模型, tester 不为空

tester: Tester | None¶: 用于模型评估的验证模型类

to_var(x: torch.Tensor) → torch.Tensor[源代码]¶

将 x 转移到对应的设备上。

参数:: x (torch.Tensor) – 数据
返回:: 张量
返回类型:: torch.Tensor

train_one_step(data: dict[str, Union[str, dgl.DGLGraph, torch.Tensor]]) → float[源代码]¶

根据 data_loader 生成的 1 批次（batch） data 将模型训练 1 步。

参数:: data (dict[str, Union[dgl.DGLGraph, torch.Tensor]]) – 训练数据
返回:: 损失值
返回类型:: float

use_early_stopping: bool¶: 是否启用早停，需要 tester 和 save_path 不为空

use_gpu: bool¶: 是否使用 gpu

valid_interval: int | None¶: 训练几轮在验证集上评估一次模型, tester 不为空

wandb_logger: WandbLogger¶: unike.utils.WandbLogger 对象

Trainer¶

Docs