torch4keras.model module

class torch4keras.model.AccelerateTrainer(module: torch.nn.Module, **configs)

accelerate来训练

compile(*args, **kwargs)

complile: 定义loss, optimizer, metrics等参数

Parameters:

• loss – loss

• optimizer – 优化器

• scheduler – scheduler

• clip_grad_norm – bool, 是否使用梯度裁剪, 默认为False

• mixed_precision – bool, 是否使用混合精度，默认为False

• metrics – str/List[str]/dict, 训练过程中需要打印的指标, loss相关指标默认会打印, 目前支持accuracy, 也支持自定义metric，形式为{key: func}

• stateful_metrics – List[str], 不滑动平均仅进行状态记录的metric，指标抖动会更加明显

• grad_accumulation_steps – int, 梯度累积步数，默认为1

• bar – str, 使用进度条的种类，从kwargs中解析，默认为keras, 可选keras, tqdm, progressbar2

Returns:

None

loss_backward(loss)

loss.backward

prepare(*args, **kwargs)

调用acclerate的prepare，如在外面评估时候需要对dev_dataloader使用

unwrap_model()

返回nn.Module模块

class torch4keras.model.BaseModel(*args: Any, **kwargs: Any)

BaseModel, 使用继承的方式来使用

class torch4keras.model.BaseModelDDP(*args: Any, **kwargs: Any)

DistributedDataParallel模式使用多gpu的方法, 父类顺序颠倒也会出问题

class torch4keras.model.BaseModelDP(*args: Any, **kwargs: Any)

DataParallel模式使用多gpu的方法, 父类顺序颠倒也会出问题

class torch4keras.model.DeepSpeedTrainer(module, config_path)

deepspeed来训练

compile(*args, log_level='warning', inference=False, master_rank=0, **kwargs)

complile: 定义loss, optimizer, metrics等参数

Parameters:

• loss – loss

• optimizer – 优化器

• scheduler – scheduler

• clip_grad_norm – bool, 是否使用梯度裁剪, 默认为False

• mixed_precision – bool, 是否使用混合精度，默认为False

• metrics – str/List[str]/dict, 训练过程中需要打印的指标, loss相关指标默认会打印, 目前支持accuracy, 也支持自定义metric，形式为{key: func}

• stateful_metrics – List[str], 不滑动平均仅进行状态记录的metric，指标抖动会更加明显

• grad_accumulation_steps – int, 梯度累积步数，默认为1

• bar – str, 使用进度条的种类，从kwargs中解析，默认为keras, 可选keras, tqdm, progressbar2

Returns:

None

loss_backward(loss)

loss.backward

resume_from_checkpoint(*args, **kwargs)

同时加载模型、优化器、训练过程参数

Parameters:

• model_path – str, 模型文件路径

• optimizer_path – str, 优化器文件路径

• scheduler_path – str, scheduler文件路径

• step_params_path – str, 训练过程参数保存路径

save_to_checkpoint(*args, **kwargs)

同时保存模型、优化器、训练过程参数、scheduler

Parameters:

• model_path – str, 模型文件路径

• optimizer_path – str, 优化器文件路径

• scheduler_path – str, scheduler文件路径

• step_params_path – str, 训练过程参数保存路径

• mapping – dict, 模型文件的mapping

step()

参数更新

unwrap_model()

返回nn.Module模块

class torch4keras.model.Trainer(module=None)

Trainer, 传入Module实例

Parameters:

module – None/nn.Module，nn.Module()的模型实例

compile(loss, optimizer, scheduler=None, clip_grad_norm=None, mixed_precision=False, metrics=None, stateful_metrics=None, grad_accumulation_steps=1, **kwargs)

complile: 定义loss, optimizer, metrics等参数

Parameters:

• loss – loss

• optimizer – 优化器

• scheduler – scheduler

• clip_grad_norm – bool, 是否使用梯度裁剪, 默认为False

• mixed_precision – bool, 是否使用混合精度，默认为False

• metrics – str/List[str]/dict, 训练过程中需要打印的指标, loss相关指标默认会打印, 目前支持accuracy, 也支持自定义metric，形式为{key: func}

• stateful_metrics – List[str], 不滑动平均仅进行状态记录的metric，指标抖动会更加明显

• grad_accumulation_steps – int, 梯度累积步数，默认为1

• bar – str, 使用进度条的种类，从kwargs中解析，默认为keras, 可选keras, tqdm, progressbar2

Returns:

None

property device: torch.device

获取model所在的device

fit(train_dataloader, steps_per_epoch=None, epochs=1, callbacks=None, verbose=1)

模型训练

Parameters:

• train_dataloader – Dataloader, 训练数据集

• steps_per_epoch – int, 每个epoch训练的steps，默认为None表示自行计算

• epochs – int, 训练的轮次, 默认为1

• callbacks – Callback/List[Callback], 回调函数，可调用预制的Callback或者自定义，默认为None

• verbose – int, 是否打印，默认为1表示打印

Returns:

None

load_steps_params(save_path)

导入训练过程参数

Parameters:

save_path – str, 训练过程参数保存路径

load_weights(load_path, strict=True, mapping={})

加载模型权重, 支持加载权重文件list

Parameters:

• save_path – str/tuple/list, 权重加载路径

• strict – bool, torch.load()是否严格加载

• mapping – dict, 指定key的映射

loss_backward(loss)

loss.backward

predict(*inputs, **input_kwargs)

模型预测，调用forward()

print_trainable_parameters()

打印可训练的参数量

resume_from_checkpoint(model_path=None, optimizer_path=None, scheduler_path=None, step_params_path=None)

同时加载模型、优化器、训练过程参数

Parameters:

• model_path – str, 模型文件路径

• optimizer_path – str, 优化器文件路径

• scheduler_path – str, scheduler文件路径

• step_params_path – str, 训练过程参数保存路径

save_steps_params(save_path)

保存训练过程参数

Parameters:

save_path – str, 训练过程参数保存路径

save_to_checkpoint(model_path=None, optimizer_path=None, scheduler_path=None, step_params_path=None, mapping={}, verbose=0)

同时保存模型、优化器、训练过程参数、scheduler

Parameters:

• model_path – str, 模型文件路径

• optimizer_path – str, 优化器文件路径

• scheduler_path – str, scheduler文件路径

• step_params_path – str, 训练过程参数保存路径

• mapping – dict, 模型文件的mapping

save_weights(save_path, mapping={}, trainable_only=False, verbose=1)

保存模型权重

Parameters:

• save_path – str, 权重保存路径

• mapping – dict, 指定key的映射

• trainable_only – bool, 指定仅保存可训练参数

step()

参数更新

to_model_device(*inputs, **input_kwargs)

遍历并转移到model.device上

unwrap_model()

返回nn.Module模块

torch4keras.model.TrainerDDP

alias of BaseModelDDP

torch4keras.model.TrainerDP

alias of BaseModelDP

torch4keras.model.add_trainer(obj, include=None, exclude=None)

为对象添加Triner对应的方法