Trình theo dõi thí nghiệm

Có rất nhiều API theo dõi thí nghiệm hiện có, tuy nhiên việc đưa tất cả chúng hoạt động trong môi trường đa tiến trình thường khá phức tạp. Accelerate cung cấp một API theo dõi chung, có thể được sử dụng để ghi lại các thông tin hữu ích trong quá trình chạy script của bạn thông qua Accelerator.log()

Trình theo dõi tích hợp

Hiện tại Accelerate hỗ trợ tám bộ theo dõi ngay khi cài đặt:

• TensorBoard
• WandB
• Trackio
• CometML
• Mục tiêu
• MLFlow
• ClearML
• DVCLive

Để sử dụng bất kỳ loại nào trong số đó, hãy truyền loại đã chọn vào tham số log_with trong Accelerate:

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])

Vào đầu thí nghiệm của bạn, hãy sử dụng Accelerator.init_trackers() để thiết lập dự án và có thể thêm các siêu tham số của thí nghiệm cần được ghi lại:

hps = {"num_iterations": 5, "learning_rate": 1e-2}
accelerator.init_trackers("my_project", config=hps)

Khi bạn sẵn sàng ghi log bất kỳ dữ liệu nào, hãy sử dụng Accelerator.log(). Bạn cũng có thể truyền vào một step để liên kết dữ liệu với một bước cụ thể trong vòng lặp huấn luyện.

accelerator.log({"train_loss": 1.12, "valid_loss": 0.8}, step=1)

Sau khi hoàn tất huấn luyện, hãy đảm bảo chạy Accelerator.end_training() để tất cả các trình theo dõi có thể thực hiện các chức năng kết thúc của chúng, nếu có.

accelerator.end_training()

Một ví dụ đầy đủ được trình bày bên dưới:

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()

Nếu một tracker yêu cầu một thư mục để lưu trữ dữ liệu, chẳng hạn như TensorBoard, thì hãy truyền đường dẫn thư mục vào project_dir. Tham số project_dir hữu ích khi có các cấu hình khác cần được kết hợp trong lớp dữ liệu ProjectConfiguration. Ví dụ, bạn có thể lưu dữ liệu TensorBoard vào project_dir và mọi thứ khác có thể được ghi log trong tham số logging_dir của [~utils.ProjectConfiguration:

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)

Thực hiện các bộ theo dõi tùy chỉnh

Để triển khai một tracker mới để sử dụng trong Accelerator, bạn có thể tạo một tracker mới bằng cách triển khai lớp GeneralTracker. Mỗi tracker phải triển khai ba hàm và có ba thuộc tính:

• __init__:

  • Nên lưu trữ run_name và khởi tạo API tracker của thư viện tích hợp.
  • Nếu một tracker lưu trữ dữ liệu cục bộ (chẳng hạn như TensorBoard), có thể thêm tham số logging_dir.

• store_init_configuration:

  • Nên nhận một từ điển values và lưu trữ chúng dưới dạng cấu hình thí nghiệm một lần

• log:

  • Nên nhận một từ điển values và một step, và nên ghi log chúng vào run

• name (str):

  • Tên chuỗi duy nhất cho bộ theo dõi, chẳng hạn như "wandb" cho bộ theo dõi wandb.
  • Điều này sẽ được sử dụng để tương tác với trình theo dõi này

• requires_logging_directory (bool):

  • Liệu logging_dir có cần thiết cho trình theo dõi cụ thể này và nó có sử dụng một thư mục ghi log hay không.

• tracker:

  • Điều này nên được triển khai dưới dạng một hàm @property
  • Nên trả về cơ chế theo dõi bên trong mà thư viện sử dụng, chẳng hạn như đối tượng run của wandb.

Mỗi phương thức cũng nên sử dụng lớp state.PartialState nếu logger chỉ nên được thực thi trên tiến trình chính, chẳng hạn.

Một ví dụ ngắn gọn có thể được xem bên dưới với tích hợp Weights and Biases, chỉ chứa thông tin liên quan và ghi log trên tiến trình chính:

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)

Khi bạn đã sẵn sàng để tạo đối tượng Accelerator, hãy truyền một thực thể của trình theo dõi của bạn vào Accelerator.log_with để nó được sử dụng tự động với API:

tracker = MyCustomTracker("some_run_name")
accelerator = Accelerator(log_with=tracker)

Chúng cũng có thể được kết hợp với các bộ theo dõi hiện có, bao gồm cả "all":

tracker = MyCustomTracker("some_run_name")
accelerator = Accelerator(log_with=[tracker, "all"])

Đăng ký bộ theo dõi tùy chỉnh theo tên

Thay vì truyền một instance, bạn có thể đăng ký một lớp tracker tùy chỉnh để nó có thể được chọn theo name, giống như các tracker tích hợp sẵn. Cách này hữu ích khi các tracker cần sử dụng được đọc từ tệp cấu hình hoặc dòng lệnh:

from accelerate import Accelerator
from accelerate.tracking import GeneralTracker, register_tracker_class


class MyCustomTracker(GeneralTracker):
    name = "my_tracker"
    requires_logging_directory = False
    # ... implement the rest of the `GeneralTracker` interface


register_tracker_class(MyCustomTracker)
accelerator = Accelerator(log_with="my_tracker")

Truy cập trình theo dõi bên trong

Nếu bạn muốn tương tác trực tiếp với một tracker tùy chỉnh, bạn có thể truy cập nhanh một tracker bằng phương thức Accelerator.get_tracker(). Chỉ cần truyền vào chuỗi tương ứng với thuộc tính .name của tracker và nó sẽ trả về tracker đó trên tiến trình chính.

Ví dụ này cho thấy cách thực hiện điều đó bằng wandb:

wandb_tracker = accelerator.get_tracker("wandb")

Từ đó, bạn có thể tương tác với đối tượng run của wandb như bình thường:

wandb_tracker.log_artifact(some_artifact_to_log)

Trackers built in Accelerate will automatically execute on the correct process, so if a tracker is only meant to be ran on the main process it will do so automatically.

Nếu bạn muốn loại bỏ hoàn toàn việc bọc của Accelerate, bạn có thể đạt được kết quả tương tự bằng cách:

wandb_tracker = accelerator.get_tracker("wandb", unwrap=True)
if accelerator.is_main_process:
    wandb_tracker.log_artifact(some_artifact_to_log)

Khi wrapper không thể hoạt động

Nếu một thư viện có API không tuân theo một .log nghiêm ngặt với một từ điển tổng thể như Neptune.AI, việc ghi log có thể được thực hiện thủ công dưới một câu lệnh 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)