Spaces:
Sleeping
Sleeping
| import os | |
| import time | |
| import torch | |
| import torch.nn as nn | |
| from torch.optim import Adam | |
| from torch.optim.lr_scheduler import ReduceLROnPlateau | |
| from tqdm import tqdm | |
| from model import get_model, count_parameters | |
| from dataset import get_dataloaders, decode_prediction, ALPHABET | |
| def greedy_decode(logits: torch.Tensor) -> list[str]: | |
| preds = logits.argmax(dim=2) | |
| preds = preds.permute(1, 0) | |
| results = [] | |
| for seq in preds: | |
| results.append(decode_prediction(seq.tolist())) | |
| return results | |
| def compute_accuracy(predictions: list[str], targets: list[str]) -> float: | |
| correct = sum(p == t for p, t in zip(predictions, targets)) | |
| return correct / len(targets) if targets else 0.0 | |
| def train_one_epoch(model, loader, optimizer, criterion, device): | |
| model.train() | |
| total_loss = 0.0 | |
| progress = tqdm(loader, desc="Обучение") | |
| for batch in progress: | |
| images = batch["images"].to(device) | |
| labels = batch["labels"] # Оставляем на CPU для быстрой склейки | |
| label_lengths = batch["label_lengths"] | |
| optimizer.zero_grad() | |
| logits = model(images) # [T, B, C] | |
| log_probs = logits.log_softmax(2) | |
| T, B, _ = logits.shape | |
| input_lengths = torch.full((B,), T, dtype=torch.long, device=device) | |
| # Быстрая склейка таргетов на CPU, затем один перенос на GPU | |
| targets_flat = torch.cat([ | |
| labels[i, :label_lengths[i]] for i in range(B) | |
| ]).to(device) | |
| label_lengths = label_lengths.to(device) | |
| loss = criterion(log_probs, targets_flat, input_lengths, label_lengths) | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0) | |
| optimizer.step() | |
| total_loss += loss.item() | |
| progress.set_postfix(loss=f"{loss.item():.4f}") | |
| return total_loss / len(loader) | |
| def validate(model, loader, criterion, device): | |
| model.eval() | |
| total_loss = 0.0 | |
| all_preds, all_targets = [], [] | |
| for batch in tqdm(loader, desc="Валидация"): | |
| images = batch["images"].to(device) | |
| labels = batch["labels"] | |
| label_lengths = batch["label_lengths"] | |
| label_strings = batch["label_strings"] | |
| logits = model(images) | |
| log_probs = logits.log_softmax(2) | |
| T, B, _ = logits.shape | |
| input_lengths = torch.full((B,), T, dtype=torch.long, device=device) | |
| targets_flat = torch.cat([ | |
| labels[i, :label_lengths[i]] for i in range(B) | |
| ]).to(device) | |
| label_lengths = label_lengths.to(device) | |
| loss = criterion(log_probs, targets_flat, input_lengths, label_lengths) | |
| total_loss += loss.item() | |
| preds = greedy_decode(logits) | |
| all_preds.extend(preds) | |
| all_targets.extend(label_strings) | |
| avg_loss = total_loss / len(loader) | |
| accuracy = compute_accuracy(all_preds, all_targets) | |
| return avg_loss, accuracy, list(zip(all_preds[:5], all_targets[:5])) | |
| def train( | |
| data_dir: str = "/mnt/d/datasetsss/mjsynth", | |
| epochs: int = 20, | |
| batch_size: int = 128, | |
| learning_rate: float = 1e-3, | |
| max_train_samples: int = None, | |
| save_dir: str = "../models", | |
| resume_from: str = None, | |
| ): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Устройство: {device}") | |
| os.makedirs(save_dir, exist_ok=True) | |
| print("\nЗагрузка данных...") | |
| train_loader, val_loader = get_dataloaders( | |
| data_dir=data_dir, | |
| batch_size=batch_size, | |
| num_workers=4, | |
| max_train_samples=max_train_samples, | |
| ) | |
| num_classes = len(ALPHABET) + 1 # +1 для BLANK | |
| model = get_model(num_classes=num_classes, device=device) | |
| print(f"Параметров в модели: {count_parameters(model):,}") | |
| criterion = nn.CTCLoss(blank=0, reduction="mean", zero_infinity=True) | |
| optimizer = Adam(model.parameters(), lr=learning_rate) | |
| scheduler = ReduceLROnPlateau(optimizer, mode="max", factor=0.5, patience=2) | |
| start_epoch = 0 | |
| best_accuracy = 0.0 | |
| if resume_from and os.path.exists(resume_from): | |
| print(f"\nЗагружаем чекпоинт: {resume_from}") | |
| checkpoint = torch.load(resume_from, map_location=device) | |
| model.load_state_dict(checkpoint["model_state_dict"]) | |
| optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) | |
| if "scheduler_state_dict" in checkpoint: | |
| scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) | |
| start_epoch = checkpoint.get("epoch", 0) | |
| best_accuracy = checkpoint.get("accuracy", 0.0) | |
| print(f"Продолжаем с эпохи {start_epoch}, лучшая точность: {best_accuracy:.4f}") | |
| print(f"\nНачинаем обучение: {epochs} эпох, batch_size={batch_size}\n") | |
| for epoch in range(start_epoch, epochs): | |
| start_time = time.time() | |
| print(f"\n=== Эпоха {epoch + 1}/{epochs} ===") | |
| train_loss = train_one_epoch(model, train_loader, optimizer, criterion, device) | |
| val_loss, val_accuracy, examples = validate(model, val_loader, criterion, device) | |
| scheduler.step(val_accuracy) | |
| elapsed = time.time() - start_time | |
| print(f"\nTrain Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Val Accuracy: {val_accuracy:.4f}") | |
| print(f"Время эпохи: {elapsed:.1f} сек") | |
| print("Примеры предсказаний (предсказано → правильно):") | |
| for pred, target in examples: | |
| mark = "✓" if pred == target else "✗" | |
| print(f" {mark} '{pred}' → '{target}'") | |
| checkpoint = { | |
| "epoch": epoch + 1, | |
| "model_state_dict": model.state_dict(), | |
| "optimizer_state_dict": optimizer.state_dict(), | |
| "scheduler_state_dict": scheduler.state_dict(), | |
| "accuracy": val_accuracy, | |
| } | |
| torch.save(checkpoint, os.path.join(save_dir, "last_checkpoint.pth")) | |
| if val_accuracy > best_accuracy: | |
| best_accuracy = val_accuracy | |
| torch.save(checkpoint, os.path.join(save_dir, "best_model.pth")) | |
| print(f"Новая лучшая модель сохранена! Точность: {best_accuracy:.4f}") | |
| print(f"\nОбучение завершено. Лучшая точность: {best_accuracy:.4f}") | |
| return model | |
| if __name__ == "__main__": | |
| train( | |
| data_dir="/mnt/d/datasatessss/mjsynth/mnt/ramdisk/max/90kDICT32px", | |
| epochs=20, | |
| batch_size=128, | |
| learning_rate=1e-3, | |
| max_train_samples=200_000, | |
| save_dir="../models", | |
| ) |