از بهینه محلی تا فروپاشی کامل: وقتی بهینه‌سازی به فاجعه تبدیل می‌شود

min_G max_D V(D, G) = E[log D(x)] + E[log(1 - D(G(z)))]

# به جای یک قدم ساده
d_loss = discriminator_loss(real, fake)

# Unrolled: چند قدم آینده را شبیه‌سازی کن
for _ in range(unroll_steps):
    # کپی Discriminator
    d_copy = copy_discriminator()
    # یک قدم update
    d_copy.update()
    # Generator با در نظر گرفتن این تغییر update می‌شود

class MinibatchDiscrimination(nn.Module):
    def forward(self, x):
        # محاسبه شباهت بین نمونه‌ها در batch
        distances = compute_pairwise_distances(x)
        # اگر همه خیلی شبیه هم باشند ← احتمالاً fake
        diversity_score = distances.mean()
        return torch.cat([x, diversity_score], dim=1)

from torch.nn.utils import spectral_norm

class Generator(nn.Module):
    def __init__(self):
        # به جای Conv2d معمولی
        self.conv1 = spectral_norm(nn.Conv2d(128, 256, 3))
        self.conv2 = spectral_norm(nn.Conv2d(256, 512, 3))

# Epoch 1-10: تصاویر 4x4
# Epoch 11-20: تصاویر 8x8  
# Epoch 21-30: تصاویر 16x16
# ...
# Epoch 61-70: تصاویر 1024x1024

# به جای Generator(noise)
output = Generator(noise, label)  # label = نوع خروجی مورد نظر

# مثال: برای تولید رقم 5
output = Generator(noise, label=5)

class ModeCollapseDetector:
    def __init__(self, threshold=0.1, window=100):
        self.threshold = threshold
        self.window = window
        self.outputs_history = []
        
    def check(self, generated_batch):
        """
        بررسی اینکه آیا Mode Collapse رخ داده
        """
        # محاسبه تنوع در batch
        # روش 1: Standard deviation of features
        features = extract_features(generated_batch)
        diversity = features.std(dim=0).mean()
        
        self.outputs_history.append(diversity)
        
        if len(self.outputs_history) > self.window:
            self.outputs_history.pop(0)
            
        # اگر تنوع کم باشد
        if diversity < self.threshold:
            recent_avg = sum(self.outputs_history) / len(self.outputs_history)
            if recent_avg < self.threshold * 1.5:
                return True, "🚨 MODE COLLAPSE DETECTED!"
        
        return False, f"Diversity: {diversity:.4f}"

# استفاده
detector = ModeCollapseDetector()
for epoch in range(num_epochs):
    fake_images = generator(noise)
    
    collapsed, msg = detector.check(fake_images)
    print(msg)
    
    if collapsed:
        # اقدامات اضطراری
        # 1. کاهش learning rate
        for param_group in g_optimizer.param_groups:
            param_group['lr'] *= 0.5
        
        # 2. افزایش noise
        noise_std *= 1.5
        
        # 3. یا restart با وزن‌های جدید
        # generator.apply(weights_init)

وزن‌های شبکه: W

Task 1: W به سمت بهینه Task 1 حرکت می‌کند ← W₁
Task 2: W₁ به سمت بهینه Task 2 حرکت می‌کند ← W₂

اما: W₂ ممکن است برای Task 1 خیلی بد باشد!

# دقت روی همه taskهای قبلی
avg_acc = sum(accuracy_on_task_i for i in previous_tasks) / num_previous_tasks

# چقدر فراموش کرده؟
forgetting = accuracy_after_task1 - accuracy_after_taskN
# اگر مثبت ← فراموش کرده
# اگر منفی ← حتی بهتر شده (backward transfer)

# آیا task قبلی به task جدید کمک کرده؟
forward_transfer = accuracy_with_pretraining - accuracy_from_scratch

Loss = Loss_task_new + λ Σ F_i (θ_i - θ*_i)²

F_i = Fisher Information Matrix (اهمیت وزن i برای task قبلی)
θ*_i = وزن بهینه برای task قبلی
λ = میزان محافظت (معمولاً 1000-10000)

class EWC:
    def __init__(self, model, dataloader, lambda_=1000):
        self.model = model
        self.lambda_ = lambda_
        self.fisher = {}
        self.optimal_params = {}
        
        # محاسبه Fisher Information
        self._compute_fisher(dataloader)
        
    def _compute_fisher(self, dataloader):
        """
        محاسبه اهمیت هر وزن برای task فعلی
        """
        self.model.eval()
        
        for name, param in self.model.named_parameters():
            self.fisher[name] = torch.zeros_like(param)
            self.optimal_params[name] = param.data.clone()
        
        for data, target in dataloader:
            self.model.zero_grad()
            output = self.model(data)
            loss = F.cross_entropy(output, target)
            loss.backward()
            
            # Fisher = گرادیان²
            for name, param in self.model.named_parameters():
                self.fisher[name] += param.grad.data ** 2
        
        # نرمال‌سازی
        for name in self.fisher:
            self.fisher[name] /= len(dataloader)
    
    def penalty(self):
        """
        جریمه برای تغییر وزن‌های مهم
        """
        loss = 0
        for name, param in self.model.named_parameters():
            fisher = self.fisher[name]
            optimal = self.optimal_params[name]
            loss += (fisher * (param - optimal) ** 2).sum()
        
        return self.lambda_ * loss

# استفاده
ewc = None
for task_id, task_data in enumerate(tasks):
    for epoch in range(num_epochs):
        for data, target in task_data:
            output = model(data)
            loss = F.cross_entropy(output, target)
            
            # اضافه کردن EWC penalty
            if ewc is not None:
                loss += ewc.penalty()
            
            loss.backward()
            optimizer.step()
    
    # بعد از هر task، EWC را update کن
    ewc = EWC(model, task_data, lambda_=5000)

class ProgressiveNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.columns = nn.ModuleList()  # هر task یک column
        self.lateral_connections = nn.ModuleList()
        
    def add_task(self, input_size, hidden_size, output_size):
        """
        اضافه کردن ستون جدید برای task جدید
        """
        new_column = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, output_size)
        )
        
        # اتصالات lateral از column‌های قبلی
        if len(self.columns) > 0:
            lateral = nn.ModuleList([
                nn.Linear(hidden_size, hidden_size)
                for _ in range(len(self.columns))
            ])
            self.lateral_connections.append(lateral)
        
        self.columns.append(new_column)
        
        # freeze کردن column‌های قبلی
        for i in range(len(self.columns) - 1):
            for param in self.columns[i].parameters():
                param.requires_grad = False
    
    def forward(self, x, task_id):
        """
        forward برای task مشخص
        """
        # column فعلی
        h = self.columns[task_id](x)
        
        # اضافه کردن اطلاعات از column‌های قبلی
        if task_id > 0:
            for i in range(task_id):
                h_prev = self.columns[i](x)
                h += self.lateral_connections[task_id-1][i](h_prev)
        
        return h

class GEM:
    def __init__(self, model, memory_size_per_task=100):
        self.model = model
        self.memory = {}  # {task_id: (data, labels)}
        self.memory_size = memory_size_per_task
        
    def store_samples(self, task_id, dataloader):
        """
        ذخیره نمونه‌های معرف از task
        """
        data_list, label_list = [], []
        
        for data, labels in dataloader:
            data_list.append(data)
            label_list.append(labels)
            if len(data_list) * data.size(0) >= self.memory_size:
                break
        
        self.memory[task_id] = (
            torch.cat(data_list)[:self.memory_size],
            torch.cat(label_list)[:self.memory_size]
        )
    
    def compute_gradient(self, task_id):
        """
        محاسبه gradient روی memory task
        """
        if task_id not in self.memory:
            return None
            
        data, labels = self.memory[task_id]
        output = self.model(data)
        loss = F.cross_entropy(output, labels)
        
        grads = torch.autograd.grad(loss, self.model.parameters())
        return grads
    
    def project_gradient(self, current_grad):
        """
        اگر gradient روی task‌های قبلی منفی است، project کن
        """
        for task_id in self.memory.keys():
            mem_grad = self.compute_gradient(task_id)
            
            # محاسبه dot product
            dot = sum((g1 * g2).sum() for g1, g2 in zip(current_grad, mem_grad))
            
            # اگر منفی است (یعنی به task قبلی آسیب می‌زند)
            if dot < 0:
                # project کن
                # g = g - ((g · m) / ||m||²) * m
                mem_norm = sum((g ** 2).sum() for g in mem_grad)
                for i, (g, m) in enumerate(zip(current_grad, mem_grad)):
                    current_grad[i] = g - (dot / mem_norm) * m
        
        return current_grad

# استفاده
gem = GEM(model, memory_size_per_task=200)
for task_id, task_data in enumerate(tasks):
    for epoch in range(num_epochs):
        for data, target in task_data:
            optimizer.zero_grad()
            output = model(data)
            loss = F.cross_entropy(output, target)
            loss.backward()
            
            # پروجکت کردن gradient
            current_grads = [p.grad.clone() for p in model.parameters()]
            projected_grads = gem.project_gradient(current_grads)
            
            # جایگزین کردن gradient‌ها
            for p, g in zip(model.parameters(), projected_grads):
                p.grad = g
            
            optimizer.step()
    
    # ذخیره نمونه‌ها برای task فعلی
    gem.store_samples(task_id, task_data)

class LwF:  # Learning without Forgetting
    def __init__(self, model, temperature=2.0, alpha=0.5):
        self.model = model
        self.old_model = None
        self.temperature = temperature
        self.alpha = alpha  # تعادل بین task جدید و قدیمی
        
    def before_new_task(self):
        """
        قبل از شروع task جدید، کپی از مدل بگیر
        """
        self.old_model = copy.deepcopy(self.model)
        self.old_model.eval()
        for param in self.old_model.parameters():
            param.requires_grad = False
    
    def distillation_loss(self, logits_new, logits_old):
        """
        محاسبه Knowledge Distillation Loss
        """
        T = self.temperature
        
        # Soft targets از مدل قدیمی
        soft_targets = F.softmax(logits_old / T, dim=1)
        # Soft predictions از مدل جدید
        soft_pred = F.log_softmax(logits_new / T, dim=1)
        
        # KL divergence
        kd_loss = F.kl_div(soft_pred, soft_targets, reduction='batchsize_average')
        kd_loss = kd_loss * (T ** 2)
        
        return kd_loss
    
    def train_step(self, data, target, new_task_data=True):
        """
        یک قدم آموزش
        """
        output = self.model(data)
        
        # Loss task جدید
        ce_loss = F.cross_entropy(output, target)
        
        total_loss = ce_loss
        
        # اگر task جدید است و مدل قدیمی داریم
        if new_task_data and self.old_model is not None:
            with torch.no_grad():
                old_output = self.old_model(data)
            
            kd_loss = self.distillation_loss(output, old_output)
            total_loss = self.alpha * ce_loss + (1 - self.alpha) * kd_loss
        
        return total_loss

# استفاده
lwf = LwF(model, temperature=2.0, alpha=0.7)
for task_id, task_data in enumerate(tasks):
    # قبل از task جدید
    if task_id > 0:
        lwf.before_new_task()
    
    for epoch in range(num_epochs):
        for data, target in task_data:
            optimizer.zero_grad()
            loss = lwf.train_step(data, target, new_task_data=True)
            loss.backward()
            optimizer.step()

gradient_layer_1 = gradient_output × W_n × W_(n-1) × ... × W_2 × W_1

اگر هر W > 1:
gradient خیلی بزرگ می‌شود (مثلاً 1.1^100 = 13780)

اگر هر W < 1:
gradient خیلی کوچک می‌شود (مثلاً 0.9^100 = 0.0000266)

# فرض کنید یک LSTM برای ترجمه ماشینی داریم
model = LSTM(input_size=1000, hidden_size=512, num_layers=4)
optimizer = Adam(model.parameters(), lr=0.001)  # LR خیلی بالا!

for epoch in range(100):
    for batch in dataloader:
        output = model(batch)
        loss = criterion(output, target)
        
        print(f"Loss: {loss.item():.4f}")
        
        loss.backward()
        optimizer.step()
# خروجی:
# Epoch 1: Loss: 4.5234
# Epoch 1: Loss: 4.1823
# Epoch 1: Loss: 3.9234
# ...
# Epoch 3: Loss: 2.1234
# Epoch 3: Loss: NaN  ← 💥 EXPLOSION!

# روش 1: Clip by norm (توصیه می‌شود)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# روش 2: Clip by value
torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)
# استفاده کامل
for epoch in range(num_epochs):
    for data, target in dataloader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        
        # 🔧 این خط فاجعه را جلوگیری می‌کند!
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        optimizer.step()

# فرض کنید gradient‌ها این‌اند:
gradients = [10.0, 50.0, 100.0, 5.0]
norm = sqrt(10² + 50² + 100² + 5²) = 112.36

# max_norm = 1.0
# اگر norm > max_norm:
scale = max_norm / norm = 1.0 / 112.36 = 0.0089
# gradient‌های جدید:
clipped_gradients = [g * scale for g in gradients]
# = [0.089, 0.445, 0.89, 0.0445]

# شروع با LR کوچک
optimizer = Adam(model.parameters(), lr=1e-4)  # نه 1e-3!

# یا استفاده از scheduler
scheduler = ReduceLROnPlateau(optimizer, factor=0.5, patience=5)
for epoch in range(num_epochs):
    train_loss = train_epoch()
    val_loss = validate()
    
    # اگر loss منفجر شد، LR را کاهش بده
    scheduler.step(val_loss)

class StableRNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size)
        # اضافه کردن BatchNorm بعد از LSTM
        self.bn = nn.BatchNorm1d(hidden_size)
        self.fc = nn.Linear(hidden_size, output_size)
    
    def forward(self, x):
        out, _ = self.lstm(x)
        out = self.bn(out)  # 🔧 Stabilizer!
        out = self.fc(out)
        return out

def init_weights(m):
    if isinstance(m, nn.Linear):
        # Xavier initialization برای جلوگیری از explosion
        nn.init.xavier_uniform_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.LSTM):
        # Orthogonal initialization برای RNN
        for name, param in m.named_parameters():
            if 'weight_hh' in name:
                nn.init.orthogonal_(param)
            elif 'weight_ih' in name:
                nn.init.xavier_uniform_(param)
            elif 'bias' in name:
                nn.init.zeros_(param)

model.apply(init_weights)

class GradientMonitor:
    def __init__(self, alert_threshold=10.0):
        self.alert_threshold = alert_threshold
        self.history = []
    
    def check_gradients(self, model):
        total_norm = 0
        for p in model.parameters():
            if p.grad is not None:
                param_norm = p.grad.data.norm(2)
                total_norm += param_norm.item() ** 2
        
        total_norm = total_norm ** 0.5
        self.history.append(total_norm)
        
        if total_norm > self.alert_threshold:
            print(f"⚠️ WARNING: Gradient norm = {total_norm:.2f}")
            return True
        
        if len(self.history) > 10:
            recent_avg = sum(self.history[-10:]) / 10
            if total_norm > recent_avg * 5:
                print(f"🚨 ALERT: Sudden spike in gradients!")
                return True
        
        return False

# استفاده
monitor = GradientMonitor(alert_threshold=5.0)
for epoch in range(num_epochs):
    for data, target in dataloader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        
        # بررسی gradient‌ها
        if monitor.check_gradients(model):
            # اقدامات اضطراری
            print("Taking emergency action...")
            # کاهش learning rate
            for param_group in optimizer.param_groups:
                param_group['lr'] *= 0.1
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

class SafeTrainer:
    def __init__(self, model, optimizer, max_grad_norm=1.0):
        self.model = model
        self.optimizer = optimizer
        self.max_grad_norm = max_grad_norm
        self.explosion_count = 0
        self.checkpoint = None
        
    def save_checkpoint(self):
        """
        ذخیره checkpoint برای recovery
        """
        self.checkpoint = {
            'model': copy.deepcopy(self.model.state_dict()),
            'optimizer': copy.deepcopy(self.optimizer.state_dict())
        }
    
    def restore_checkpoint(self):
        """
        بازگردانی به checkpoint قبلی
        """
        if self.checkpoint is not None:
            self.model.load_state_dict(self.checkpoint['model'])
            self.optimizer.load_state_dict(self.checkpoint['optimizer'])
            print("✅ Restored from checkpoint")
    
    def train_step(self, data, target, criterion):
        """
        یک قدم آموزش امن
        """
        self.optimizer.zero_grad()
        
        # Forward
        output = self.model(data)
        loss = criterion(output, target)
        
        # بررسی loss
        if torch.isnan(loss) or torch.isinf(loss):
            print("🚨 NaN/Inf detected in loss!")
            self.explosion_count += 1
            
            if self.explosion_count > 3:
                print("💥 Too many explosions! Restoring checkpoint...")
                self.restore_checkpoint()
                self.explosion_count = 0
            
            return None
        
        # Backward
        loss.backward()
        
        # محاسبه gradient norm قبل از clipping
        grad_norm = 0
        for p in self.model.parameters():
            if p.grad is not None:
                grad_norm += p.grad.data.norm(2).item() ** 2
        grad_norm = grad_norm ** 0.5
        
        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(
            self.model.parameters(), 
            max_norm=self.max_grad_norm
        )
        
        # Update
        self.optimizer.step()
        
        # ذخیره checkpoint هر 100 step
        if hasattr(self, 'step_count'):
            self.step_count += 1
            if self.step_count % 100 == 0:
                self.save_checkpoint()
        else:
            self.step_count = 1
        
        return {
            'loss': loss.item(),
            'grad_norm': grad_norm
        }

# استفاده
safe_trainer = SafeTrainer(model, optimizer, max_grad_norm=1.0)
for epoch in range(num_epochs):
    for data, target in dataloader:
        result = safe_trainer.train_step(data, target, criterion)
        
        if result is not None:
            print(f"Loss: {result['loss']:.4f}, Grad Norm: {result['grad_norm']:.4f}")

فرض کنید شبکه 100 لایه دارد
هر لایه: activation = sigmoid(Wx + b)

gradient در لایه 100 = 1.0
gradient در لایه 50 = 0.01
gradient در لایه 10 = 0.0000001  ← تقریباً صفر!
gradient در لایه 1 = 10^-20  ← کاملاً صفر!

# Sigmoid
def sigmoid(x):
    return 1 / (1 + exp(-x))

# مشتق sigmoid:
d_sigmoid = sigmoid(x) * (1 - sigmoid(x))
# برای x = 10:  d_sigmoid ≈ 0.00005
# برای x = -10: d_sigmoid ≈ 0.00005
# در نواحی اشباع، مشتق تقریباً صفر!

# اگر وزن‌ها خیلی کوچک باشند:
W = 0.01 * np.random.randn(1000, 1000)

# بعد از 10 لایه:
# signal = W^10 ≈ (0.01)^10 = 10^-20  ← محو شدن!

# Pre-ResNet: شبکه‌های 50+ لایه
# Vanishing gradient باعث می‌شد فقط لایه‌های آخر یاد بگیرند
# پارادوکس: شبکه عمیق‌تر = عملکرد بدتر!

# ❌ Bad: Sigmoid (مشتق کوچک)
activation = nn.Sigmoid()

# ✅ Good: ReLU (مشتق 1 برای x > 0)
activation = nn.ReLU()
# ✅ Better: Leaky ReLU (مشتق غیرصفر برای همه x)
activation = nn.LeakyReLU(negative_slope=0.01)
# ✅ Best for Transformers: GELU
activation = nn.GELU()

Sigmoid: d/dx = σ(x)(1-σ(x))  ← max = 0.25
Tanh:    d/dx = 1 - tanh²(x)  ← max = 1.0
ReLU:    d/dx = 1 (x>0), 0 (x≤0)
Leaky:   d/dx = 1 (x>0), 0.01 (x<0)

class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.layer1 = nn.Linear(dim, dim)
        self.layer2 = nn.Linear(dim, dim)
        self.activation = nn.ReLU()
    
    def forward(self, x):
        residual = x  # ذخیره ورودی
        
        out = self.layer1(x)
        out = self.activation(out)
        out = self.layer2(out)
        
        # اضافه کردن residual
        out = out + residual  # 🔧 این خط gradient flow را نجات می‌دهد!
        out = self.activation(out)
        
        return out

بدون skip connection:
gradient_layer1 = gradient × W_100 × W_99 × ... × W_2  ← صفر می‌شود

با skip connection:
gradient_layer1 = gradient × (1 + W_100 × W_99 × ... × W_2)  ← حداقل gradient اصلی باقی می‌ماند!

class DeepNetworkWithBN(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList()
        
        for i in range(100):  # 100 لایه!
            self.layers.append(nn.Sequential(
                nn.Linear(512, 512),
                nn.BatchNorm1d(512),  # 🔧 جلوگیری از vanishing
                nn.ReLU()
            ))
    
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

# ❌ Bad: RNN ساده
class SimpleRNN(nn.Module):
    def forward(self, x, h):
        h = tanh(Wx @ x + Wh @ h + b)  # Vanishing gradient!
        return h

# ✅ Good: LSTM با gate mechanism
class LSTMCell(nn.Module):
    def forward(self, x, h, c):
        # Gate‌ها اجازه می‌دهند gradient مستقیماً flow کند
        f = sigmoid(Wf @ [x, h] + bf)  # Forget gate
        i = sigmoid(Wi @ [x, h] + bi)  # Input gate
        o = sigmoid(Wo @ [x, h] + bo)  # Output gate
        
        c_new = f * c + i * tanh(Wc @ [x, h] + bc)
        h_new = o * tanh(c_new)
        
        return h_new, c_new

from torch.utils.checkpoint import checkpoint

class VeryDeepNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([
            ResidualBlock(512) for _ in range(1000)
        ])
    
    def forward(self, x):
        # استفاده از checkpointing برای صرفه‌جویی حافظه
        for layer in self.layers:
            x = checkpoint(layer, x)  # فقط activation‌های لازم را نگه می‌دارد
        return x

# مثال loss نوسانی:
Epoch 1: Loss = 2.5
Epoch 2: Loss = 2.1
Epoch 3: Loss = 1.8
Epoch 4: Loss = 3.2  ← 💥 چرا بدتر شد؟
Epoch 5: Loss = 1.5
Epoch 6: Loss = 2.9  ← 💥 دوباره!
Epoch 7: Loss = 1.2

# LR = 0.1 برای شبکه عمیق ← خیلی بالا!
# مدل از minimum می‌پرد به یک طرف و آن طرف
optimizer = Adam(model.parameters(), lr=0.1)  # ❌
optimizer = Adam(model.parameters(), lr=1e-4)  # ✅

# Batch size = 2 ← Gradient خیلی نویزی
dataloader = DataLoader(dataset, batch_size=2)  # ❌
dataloader = DataLoader(dataset, batch_size=32)  # ✅

# اگر بعضی label‌ها اشتباه باشند
# مدل confused می‌شود و نوسان می‌کند

class WarmupScheduler:
    def __init__(self, optimizer, warmup_steps, total_steps):
        self.optimizer = optimizer
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.step_count = 0
        self.base_lr = optimizer.param_groups[0]['lr']
    
    def step(self):
        self.step_count += 1
        
        if self.step_count < self.warmup_steps:
            # Warmup: افزایش تدریجی
            lr = self.base_lr * (self.step_count / self.warmup_steps)
        else:
            # Decay: کاهش تدریجی
            progress = (self.step_count - self.warmup_steps) / (self.total_steps - self.warmup_steps)
            lr = self.base_lr * (1 - progress)
        
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

# استفاده
scheduler = WarmupScheduler(optimizer, warmup_steps=1000, total_steps=10000)
for epoch in range(num_epochs):
    for data, target in dataloader:
        # آموزش
        loss = train_step(data, target)
        
        # Update learning rate
        scheduler.step()

accumulation_steps = 4  # هر 4 step یکبار update

optimizer.zero_grad()
for i, (data, target) in enumerate(dataloader):
    output = model(data)
    loss = criterion(output, target)
    
    # تقسیم loss بر accumulation steps
    loss = loss / accumulation_steps
    loss.backward()
    
    # هر accumulation_steps یکبار update
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
# نتیجه: مثل batch size 4 برابر بزرگتر!

class EMA:
    def __init__(self, model, decay=0.999):
        self.model = model
        self.decay = decay
        self.shadow = {}
        self.backup = {}
        
        # ذخیره کپی از وزن‌ها
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()
    
    def update(self):
        """
        Update EMA weights
        """
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]
                self.shadow[name] = new_average.clone()
    
    def apply_shadow(self):
        """
        استفاده از EMA weights برای inference
        """
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.backup[name] = param.data.clone()
                param.data = self.shadow[name]
    
    def restore(self):
        """
        برگرداندن وزن‌های اصلی
        """
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                param.data = self.backup[name]
        self.backup = {}

# استفاده
ema = EMA(model, decay=0.9999)
for epoch in range(num_epochs):
    for data, target in dataloader:
        # آموزش عادی
        loss = train_step(data, target)
        
        # Update EMA
        ema.update()
    
    # Evaluation با EMA weights
    ema.apply_shadow()
    val_accuracy = evaluate(val_loader)
    ema.restore()
    
    print(f"Validation Accuracy: {val_accuracy:.2f}%")