Accelerate 文档
实验跟踪器
并获得增强的文档体验
开始使用
实验跟踪器
有大量的实验跟踪 API 可用,但让它们在多处理环境中协同工作通常很复杂。Accelerate 提供了一个通用的跟踪 API,可以通过 Accelerator.log()
在脚本中记录有用的项目。
集成跟踪器
目前,Accelerate
开箱即用地支持八种跟踪器:
- TensorBoard
- WandB
- Trackio
- CometML
- Aim
- MLFlow
- ClearML
- DVCLive
要使用它们中的任何一个,只需将所选类型传递给 Accelerate
的 log_with
参数。
from accelerate import Accelerator
from accelerate.utils import LoggerType
accelerator = Accelerator(log_with="all") # For all available trackers in the environment
accelerator = Accelerator(log_with="wandb")
accelerator = Accelerator(log_with=["wandb", LoggerType.TENSORBOARD])
在实验开始时,应使用 Accelerator.init_trackers()
来设置你的项目,并可能添加任何要记录的实验超参数。
hps = {"num_iterations": 5, "learning_rate": 1e-2}
accelerator.init_trackers("my_project", config=hps)
当你准备好记录任何数据时,应使用 Accelerator.log()
。也可以传入一个 step
,将数据与训练循环中的特定步骤关联起来。
accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)
训练完成后,请确保运行 Accelerator.end_training(),以便所有跟踪器可以运行它们的收尾功能(如果它们有的话)。
accelerator.end_training()
下面是一个完整的示例。
from accelerate import Accelerator
accelerator = Accelerator(log_with="all")
config = {
"num_iterations": 5,
"learning_rate": 1e-2,
"loss_function": str(my_loss_function),
}
accelerator.init_trackers("example_project", config=config)
my_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, my_optimizer, my_training_dataloader)
device = accelerator.device
my_model.to(device)
for iteration in range(config["num_iterations"]):
for step, batch in enumerate(my_training_dataloader):
my_optimizer.zero_grad()
inputs, targets = batch
inputs = inputs.to(device)
targets = targets.to(device)
outputs = my_model(inputs)
loss = my_loss_function(outputs, targets)
accelerator.backward(loss)
my_optimizer.step()
accelerator.log({"training_loss": loss}, step=step)
accelerator.end_training()
如果一个跟踪器需要一个目录来保存数据,例如 TensorBoard
,则将目录路径传递给 project_dir
。当有其他配置需要与 ProjectConfiguration 数据类中的配置结合时,project_dir
参数非常有用。例如,你可以将 TensorBoard 数据保存到 project_dir
,而其他所有内容都可以记录在 [~utils.ProjectConfiguration
的 logging_dir
参数中。
accelerator = Accelerator(log_with="tensorboard", project_dir=".")
# use with ProjectConfiguration
config = ProjectConfiguration(project_dir=".", logging_dir="another/directory")
accelerator = Accelerator(log_with="tensorboard", project_config=config)
实现自定义跟踪器
要在 Accelerator
中实现一个新的跟踪器,可以通过实现 GeneralTracker
类来创建一个。每个跟踪器必须实现三个函数并具有三个属性。
__init__
:- 应存储一个
run_name
并初始化集成库的跟踪器 API。 - 如果跟踪器在本地存储数据(如 TensorBoard),可以添加一个
logging_dir
参数。
- 应存储一个
store_init_configuration
:- 应接收一个
values
字典,并将其作为一次性的实验配置进行存储。
- 应接收一个
log
:- 应接收一个
values
字典和一个step
,并将它们记录到运行中。
- 应接收一个
name
(str
)- 跟踪器的唯一字符串名称,例如 wandb 跟踪器的
"wandb"
。 - 这将用于与此跟踪器进行特定交互。
- 跟踪器的唯一字符串名称,例如 wandb 跟踪器的
requires_logging_directory
(bool
)- 此特定跟踪器是否需要
logging_dir
以及是否使用它。
- 此特定跟踪器是否需要
tracker
:- 这应该实现为一个
@property
函数。 - 应返回库使用的内部跟踪机制,例如
wandb
的run
对象。
- 这应该实现为一个
如果记录器只应在主进程上执行,每个方法也应利用 state.PartialState 类。
下面是一个与 Weights and Biases 集成的简短示例,仅包含相关信息,并且只在主进程上进行日志记录。
from accelerate.tracking import GeneralTracker, on_main_process
from typing import Optional
import wandb
class MyCustomTracker(GeneralTracker):
name = "wandb"
requires_logging_directory = False
@on_main_process
def __init__(self, run_name: str):
self.run_name = run_name
run = wandb.init(self.run_name)
@property
def tracker(self):
return self.run.run
@on_main_process
def store_init_configuration(self, values: dict):
wandb.config(values)
@on_main_process
def log(self, values: dict, step: Optional[int] = None):
wandb.log(values, step=step)
当你准备好构建你的 Accelerator
对象时,将你的跟踪器的一个**实例**传递给 Accelerator.log_with
,它就会自动被用于 API。
tracker = MyCustomTracker("some_run_name")
accelerator = Accelerator(log_with=tracker)
这些也可以与现有的跟踪器混合使用,包括与 "all"
一起使用。
tracker = MyCustomTracker("some_run_name")
accelerator = Accelerator(log_with=[tracker, "all"])
访问内部跟踪器
如果希望直接与某个跟踪器进行一些自定义交互,可以使用 Accelerator.get_tracker() 方法快速访问它。只需传入与跟踪器 .name
属性对应的字符串,它就会在主进程上返回该跟踪器。
这个例子展示了如何用 wandb 来做到这一点。
wandb_tracker = accelerator.get_tracker("wandb")
从那里你可以像往常一样与 wandb
的 run
对象进行交互。
wandb_tracker.log_artifact(some_artifact_to_log)
如果你想完全移除 Accelerate 的包装,你可以通过以下方式达到同样的效果:
wandb_tracker = accelerator.get_tracker("wandb", unwrap=True)
if accelerator.is_main_process:
wandb_tracker.log_artifact(some_artifact_to_log)
当包装器无法工作时
如果一个库的 API 不遵循严格的 .log
和一个整体字典的模式,例如 Neptune.AI,那么可以在一个 if accelerator.is_main_process
语句下手动进行日志记录。
from accelerate import Accelerator
+ import neptune
accelerator = Accelerator()
+ run = neptune.init_run(...)
my_model, my_optimizer, my_training_dataloader = accelerate.prepare(my_model, my_optimizer, my_training_dataloader)
device = accelerator.device
my_model.to(device)
for iteration in config["num_iterations"]:
for batch in my_training_dataloader:
my_optimizer.zero_grad()
inputs, targets = batch
inputs = inputs.to(device)
targets = targets.to(device)
outputs = my_model(inputs)
loss = my_loss_function(outputs, targets)
total_loss += loss
accelerator.backward(loss)
my_optimizer.step()
+ if accelerator.is_main_process:
+ run["logs/training/batch/loss"].log(loss)