from accelerate import Accelerator
# wandb로 로깅하도록 Accelerator 객체에 지시합니다
accelerator = Accelerator(log_with="wandb")
# wandb run을 초기화하고, wandb 매개변수와 설정 정보를 전달합니다
accelerator.init_trackers(
project_name="my_project",
config={"dropout": 0.1, "learning_rate": 1e-2},
init_kwargs={"wandb": {"entity": "my-wandb-team"}}
)
...
# accelerator.log()를 호출하여 wandb에 로깅합니다. step은 선택 사항입니다
accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=global_step)
# wandb tracker가 올바르게 종료되도록 합니다
accelerator.end_training()