import numpy as np
import matplotlib.pyplot as plt

def step_decay(epoch, initial_lr=0.1, drop_rate=0.1, epochs_drop=30):
    """Step Decay"""
    return initial_lr * (drop_rate ** (epoch // epochs_drop))

def exponential_decay(epoch, initial_lr=0.1, decay_rate=0.95):
    """Exponential Decay"""
    return initial_lr * (decay_rate ** epoch)

def cosine_annealing(epoch, total_epochs=100, lr_max=0.1, lr_min=0.0):
    """Cosine Annealing"""
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + np.cos(np.pi * epoch / total_epochs))

def warmup_cosine(epoch, total_epochs=100, warmup_epochs=10, lr_max=0.1, lr_min=0.0):
    """Warmup + Cosine Annealing"""
    if epoch < warmup_epochs:
        return lr_max * epoch / warmup_epochs
    else:
        progress = (epoch - warmup_epochs) / (total_epochs - warmup_epochs)
        return lr_min + 0.5 * (lr_max - lr_min) * (1 + np.cos(np.pi * progress))

def one_cycle(epoch, total_epochs=100, lr_max=0.1, lr_min=0.001, pct_start=0.3):
    """OneCycleLR"""
    if epoch < total_epochs * pct_start:
        # 増加フェーズ
        progress = epoch / (total_epochs * pct_start)
        return lr_min + (lr_max - lr_min) * progress
    else:
        # 減少フェーズ
        progress = (epoch - total_epochs * pct_start) / (total_epochs * (1 - pct_start))
        return lr_max - (lr_max - lr_min) * progress

def cosine_warm_restarts(epoch, total_epochs=100, lr_max=0.1, lr_min=0.0, T_0=20, T_mult=2):
    """Cosine Annealing with Warm Restarts"""
    # 現在の周期を特定
    T_cur = epoch
    T_i = T_0

    while T_cur >= T_i:
        T_cur -= T_i
        T_i *= T_mult
        if T_i > total_epochs:
            T_i = total_epochs - epoch
            break

    return lr_min + 0.5 * (lr_max - lr_min) * (1 + np.cos(np.pi * T_cur / T_i))

# 可視化
epochs = np.arange(100)

plt.figure(figsize=(14, 10))

schedules = {
    'Step Decay': [step_decay(e) for e in epochs],
    'Exponential Decay': [exponential_decay(e) for e in epochs],
    'Cosine Annealing': [cosine_annealing(e) for e in epochs],
    'Warmup + Cosine': [warmup_cosine(e) for e in epochs],
    'OneCycleLR': [one_cycle(e) for e in epochs],
    'Cosine Warm Restarts': [cosine_warm_restarts(e) for e in epochs],
}

for idx, (name, lr_values) in enumerate(schedules.items(), 1):
    plt.subplot(2, 3, idx)
    plt.plot(epochs, lr_values, linewidth=2)
    plt.xlabel('Epoch')
    plt.ylabel('Learning Rate')
    plt.title(name)
    plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import (
    StepLR, ExponentialLR, CosineAnnealingLR,
    CosineAnnealingWarmRestarts, OneCycleLR, ReduceLROnPlateau
)
import matplotlib.pyplot as plt

# ダミーモデル
model = nn.Linear(10, 2)
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 各スケジューラの学習率推移を取得
def get_lr_history(scheduler_fn, n_epochs=100):
    """スケジューラの学習率履歴を取得"""
    model = nn.Linear(10, 2)
    optimizer = optim.SGD(model.parameters(), lr=0.1)
    scheduler = scheduler_fn(optimizer)

    lr_history = []
    for epoch in range(n_epochs):
        lr_history.append(optimizer.param_groups[0]['lr'])
        # ダミーの学習ステップ
        optimizer.step()
        scheduler.step()

    return lr_history

# 各スケジューラを作成
schedulers = {
    'StepLR': lambda opt: StepLR(opt, step_size=30, gamma=0.1),
    'ExponentialLR': lambda opt: ExponentialLR(opt, gamma=0.95),
    'CosineAnnealingLR': lambda opt: CosineAnnealingLR(opt, T_max=100, eta_min=0.001),
    'CosineAnnealingWarmRestarts': lambda opt: CosineAnnealingWarmRestarts(opt, T_0=20, T_mult=2),
}

# 可視化
plt.figure(figsize=(12, 8))

for idx, (name, scheduler_fn) in enumerate(schedulers.items(), 1):
    lr_history = get_lr_history(scheduler_fn)
    plt.subplot(2, 2, idx)
    plt.plot(lr_history, linewidth=2)
    plt.xlabel('Epoch')
    plt.ylabel('Learning Rate')
    plt.title(f'PyTorch {name}')
    plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

class WarmupCosineScheduler:
    """Warmup + Cosine Annealing スケジューラ"""

    def __init__(self, optimizer, warmup_epochs, total_epochs, lr_max, lr_min=0.0):
        self.optimizer = optimizer
        self.warmup_epochs = warmup_epochs
        self.total_epochs = total_epochs
        self.lr_max = lr_max
        self.lr_min = lr_min
        self.current_epoch = 0

    def step(self):
        if self.current_epoch < self.warmup_epochs:
            # Warmupフェーズ
            lr = self.lr_max * self.current_epoch / self.warmup_epochs
        else:
            # Cosine Annealingフェーズ
            progress = (self.current_epoch - self.warmup_epochs) / (self.total_epochs - self.warmup_epochs)
            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (1 + np.cos(np.pi * progress))

        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

        self.current_epoch += 1

    def get_lr(self):
        return self.optimizer.param_groups[0]['lr']

# 使用例
model = nn.Linear(10, 2)
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = WarmupCosineScheduler(
    optimizer,
    warmup_epochs=10,
    total_epochs=100,
    lr_max=0.001,
    lr_min=1e-6
)

lr_history = []
for epoch in range(100):
    lr_history.append(scheduler.get_lr())
    scheduler.step()

plt.figure(figsize=(10, 5))
plt.plot(lr_history, linewidth=2)
plt.axvline(x=10, color='red', linestyle='--', label='End of Warmup')
plt.xlabel('Epoch')
plt.ylabel('Learning Rate')
plt.title('Warmup + Cosine Annealing')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR, CosineAnnealingLR, OneCycleLR
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import matplotlib.pyplot as plt

def create_dataset(n_samples=5000, input_dim=20, num_classes=10):
    np.random.seed(42)
    X = np.random.randn(n_samples, input_dim).astype(np.float32)
    W = np.random.randn(input_dim, num_classes).astype(np.float32)
    logits = X @ W + np.random.randn(n_samples, num_classes).astype(np.float32) * 0.3
    y = np.argmax(logits, axis=1)
    split = int(0.8 * n_samples)
    return (X[:split], y[:split]), (X[split:], y[split:])

class MLP(nn.Module):
    def __init__(self, input_dim=20, hidden_dim=128, num_classes=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_classes)
        )

    def forward(self, x):
        return self.net(x)

def train_with_scheduler(scheduler_name, train_loader, test_loader, n_epochs=100):
    torch.manual_seed(42)
    model = MLP()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

    # スケジューラの選択
    if scheduler_name == 'No Schedule':
        scheduler = None
    elif scheduler_name == 'StepLR':
        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
    elif scheduler_name == 'CosineAnnealingLR':
        scheduler = CosineAnnealingLR(optimizer, T_max=n_epochs, eta_min=0.001)
    elif scheduler_name == 'OneCycleLR':
        scheduler = OneCycleLR(
            optimizer,
            max_lr=0.1,
            epochs=n_epochs,
            steps_per_epoch=len(train_loader)
        )
    else:
        raise ValueError(f"Unknown scheduler: {scheduler_name}")

    train_losses = []
    test_accuracies = []
    lr_history = []

    for epoch in range(n_epochs):
        # Training
        model.train()
        epoch_loss = 0
        for X_batch, y_batch in train_loader:
            optimizer.zero_grad()
            outputs = model(X_batch)
            loss = criterion(outputs, y_batch)
            loss.backward()
            optimizer.step()

            if scheduler_name == 'OneCycleLR':
                scheduler.step()

            epoch_loss += loss.item()

        train_losses.append(epoch_loss / len(train_loader))
        lr_history.append(optimizer.param_groups[0]['lr'])

        # スケジューラの更新（OneCycleLR以外）
        if scheduler is not None and scheduler_name != 'OneCycleLR':
            scheduler.step()

        # Evaluation
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for X_batch, y_batch in test_loader:
                outputs = model(X_batch)
                _, predicted = torch.max(outputs, 1)
                total += y_batch.size(0)
                correct += (predicted == y_batch).sum().item()
        test_accuracies.append(correct / total)

    return train_losses, test_accuracies, lr_history

# データ準備
(X_train, y_train), (X_test, y_test) = create_dataset()
train_dataset = TensorDataset(torch.tensor(X_train), torch.tensor(y_train))
test_dataset = TensorDataset(torch.tensor(X_test), torch.tensor(y_test))
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 各スケジューラで学習
schedulers = ['No Schedule', 'StepLR', 'CosineAnnealingLR', 'OneCycleLR']
results = {}

for name in schedulers:
    print(f"Training with {name}...")
    train_losses, test_accs, lr_history = train_with_scheduler(name, train_loader, test_loader)
    results[name] = {
        'train_loss': train_losses,
        'test_acc': test_accs,
        'lr_history': lr_history
    }
    print(f"  Final Test Accuracy: {test_accs[-1]:.4f}")

# 可視化
fig, axes = plt.subplots(1, 3, figsize=(15, 4))

colors = {'No Schedule': 'gray', 'StepLR': 'blue', 'CosineAnnealingLR': 'green', 'OneCycleLR': 'red'}

# 学習率
ax1 = axes[0]
for name in schedulers:
    ax1.plot(results[name]['lr_history'], label=name, color=colors[name], linewidth=1.5)
ax1.set_xlabel('Epoch')
ax1.set_ylabel('Learning Rate')
ax1.set_title('Learning Rate Schedule')
ax1.legend()
ax1.grid(True, alpha=0.3)
ax1.set_yscale('log')

# 訓練損失
ax2 = axes[1]
for name in schedulers:
    ax2.plot(results[name]['train_loss'], label=name, color=colors[name], linewidth=1.5)
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Training Loss')
ax2.set_title('Training Loss')
ax2.legend()
ax2.grid(True, alpha=0.3)

# テスト精度
ax3 = axes[2]
for name in schedulers:
    ax3.plot(results[name]['test_acc'], label=name, color=colors[name], linewidth=1.5)
ax3.set_xlabel('Epoch')
ax3.set_ylabel('Test Accuracy')
ax3.set_title('Test Accuracy')
ax3.legend()
ax3.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()