Weights & Biases

import wandb
import torch
import torch.nn as nn

# 1. W&B 초기화
wandb.init(
    project="my-ml-project",
    config={
        "learning_rate": 0.001,
        "epochs": 10,
        "batch_size": 32,
        "architecture": "ResNet18"
    }
)
config = wandb.config

# 2. 모델 정의
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

# 3. 학습 루프 with W&B 로깅
optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
criterion = nn.CrossEntropyLoss()

for epoch in range(config.epochs):
    # 더미 학습 (실제로는 DataLoader 사용)
    x = torch.randn(config.batch_size, 784)
    y = torch.randint(0, 10, (config.batch_size,))

    optimizer.zero_grad()
    output = model(x)
    loss = criterion(output, y)
    loss.backward()
    optimizer.step()

    accuracy = (output.argmax(dim=1) == y).float().mean()

    # 메트릭 로깅
    wandb.log({
        "epoch": epoch,
        "train_loss": loss.item(),
        "train_accuracy": accuracy.item(),
        "learning_rate": optimizer.param_groups[0]["lr"]
    })

    print(f"Epoch {epoch}: loss={loss.item():.4f}, acc={accuracy.item():.4f}")

# 4. 모델 저장 (Artifact)
artifact = wandb.Artifact("model", type="model")
torch.save(model.state_dict(), "model.pt")
artifact.add_file("model.pt")
wandb.log_artifact(artifact)

# 5. 하이퍼파라미터 스윕 설정 예시
sweep_config = {
    "method": "bayes",  # bayes, grid, random
    "metric": {"name": "val_loss", "goal": "minimize"},
    "parameters": {
        "learning_rate": {"min": 0.0001, "max": 0.1},
        "batch_size": {"values": [16, 32, 64]},
        "epochs": {"value": 10}
    }
}

# sweep_id = wandb.sweep(sweep_config, project="my-ml-project")
# wandb.agent(sweep_id, train_function, count=20)

# 6. Hugging Face Transformers 통합
# from transformers import Trainer, TrainingArguments
#
# training_args = TrainingArguments(
#     output_dir="./results",
#     report_to="wandb",  # W&B 자동 통합
#     run_name="bert-finetune"
# )

wandb.finish()