!pip install wandb
import wandb

# 1. 组织超参数
config = {'learning_rate': 0.001}

# 2. 启动wandb
wandb.init(project='gpt5', config=config)

# 在这里模型训练

# 3. 记录指标来可视化表现
wandb.log({"loss": loss})

# 4. 在notbook里运行时，记得结束
wandb.finish()

# 导入相应库
import math
from pathlib import Path
from types import SimpleNamespace
from tqdm.auto import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
from utilities import get_dataloaders

import wandb

INPUT_SIZE = 3 * 16 * 16
OUTPUT_SIZE = 5
HIDDEN_SIZE = 256
NUM_WORKERS = 2
CLASSES = ["hero", "non-hero", "food", "spell", "side-facing"]
DATA_DIR = Path('./data/')
DEVICE = torch.device("cuda" if torch.cuda.is_available()  else "cpu")

def get_model(dropout):
    "有 Dropout 的简单多层感知机"
    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(INPUT_SIZE, HIDDEN_SIZE),
        nn.BatchNorm1d(HIDDEN_SIZE),
        nn.ReLU(),
        nn.Dropout(dropout),
        nn.Linear(HIDDEN_SIZE, OUTPUT_SIZE)
    ).to(DEVICE)

# 定义一个配置对象来储存我们的超参数
config = SimpleNamespace(
    epochs = 2,
    batch_size = 128,
    lr = 1e-5,
    dropout = 0.5,
    slice_size = 10_000,
    valid_pct = 0.2,
)

def train_model(config):
    "用已知的配置来训练模型"
    
    # 启用 wandb 来传递项目名和配置对象
    wandb.init(
        project="dlai_intro",
        config=config,
    )

    # 获取数据
    train_dl, valid_dl = get_dataloaders(DATA_DIR, 
                                         config.batch_size, 
                                         config.slice_size, 
                                         config.valid_pct)
    n_steps_per_epoch = math.ceil(len(train_dl.dataset) / config.batch_size)

    # 一个简单的多层感知机模型
    model = get_model(config.dropout)

    # 制取损失和优化器
    loss_func = nn.CrossEntropyLoss()
    optimizer = Adam(model.parameters(), lr=config.lr)

    example_ct = 0

    for epoch in tqdm(range(config.epochs), total=config.epochs):
        model.train()

        for step, (images, labels) in enumerate(train_dl):
            images, labels = images.to(DEVICE), labels.to(DEVICE)

            outputs = model(images)
            train_loss = loss_func(outputs, labels)
            optimizer.zero_grad()
            train_loss.backward()
            optimizer.step()

            example_ct += len(images)
            metrics = {
                "train/train_loss": train_loss,
                "train/epoch": epoch + 1,
                "train/example_ct": example_ct
            }
            # 记录相应指标
            wandb.log(metrics)
            
        # 计算验证指标，记录上最近epoch的图像
        val_loss, accuracy = validate_model(model, valid_dl, loss_func)
        # 计算训练和验证指标
        val_metrics = {
            "val/val_loss": val_loss,
            "val/val_accuracy": accuracy
        }
        # 记录相应指标
        wandb.log(val_metrics)
    
    # 结束
    wandb.finish()

def validate_model(model, valid_dl, loss_func):
    "在验证集上计算模型的表现"
    model.eval()
    val_loss = 0.0
    correct = 0

    with torch.inference_mode():
        for i, (images, labels) in enumerate(valid_dl):
            images, labels = images.to(DEVICE), labels.to(DEVICE)

            # 向前传递
            outputs = model(images)
            val_loss += loss_func(outputs, labels) * labels.size(0)

            # 计算准确度并累积
            _, predicted = torch.max(outputs.data, 1)
            correct += (predicted == labels).sum().item()
            
    return val_loss / len(valid_dl.dataset), correct / len(valid_dl.dataset)

wandb.login(anonymous="allow")

train_model(config)

# 我们现在改变学习率为1e-4 
# 然后我们看看这会怎样影响我们的结果
config.lr = 1e-4
train_model(config)

config.dropout = 0.1
config.epochs = 1
train_model(config)

config.lr = 1e-3
train_model(config)