Spaces:
Sleeping
Sleeping
File size: 5,311 Bytes
4af0d61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | import torch
from torch.utils.data import DataLoader
from datasets import load_dataset
from torchvision import transforms
from huggingface_hub import HfApi, create_repo, upload_folder
import os
from tqdm.auto import tqdm
import json
class LouTrainer:
def __init__(self, model, diffusion, device='cuda'):
self.model = model
self.diffusion = diffusion
self.device = device
self.optimizer = None
self.current_epoch = 0
def setup_optimizer(self, lr=1e-4):
self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=0.01)
self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=1000)
def prepare_dataset(self, dataset_name="diffusers/pokemon-gpt4-captions", batch_size=4, image_size=256):
"""Charge et prépare le dataset"""
# Transformations
transform = transforms.Compose([
transforms.Resize(image_size),
transforms.CenterCrop(image_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
# Chargement dataset
dataset = load_dataset(dataset_name, split="train")
def transform_examples(examples):
images = [transform(img.convert("RGB")) for img in examples["image"]]
return {"pixel_values": images}
dataset.set_transform(transform_examples)
dataloader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
num_workers=0,
pin_memory=True
)
return dataloader
def train_epoch(self, dataloader, progress_bar=None):
"""Entraîne une epoch"""
total_loss = 0
for batch_idx, batch in enumerate(dataloader):
images = batch["pixel_values"].to(self.device)
self.optimizer.zero_grad()
loss = self.diffusion.training_step(images)
loss.backward()
# Gradient clipping
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
self.optimizer.step()
total_loss += loss.item()
if progress_bar:
progress_bar.update(1)
progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})
self.scheduler.step()
return total_loss / len(dataloader)
def train(self, dataset_name, num_epochs, lr, batch_size=4, image_size=256,
save_every=5, repo_id=None, token=None, progress_callback=None):
"""Entraînement complet"""
self.setup_optimizer(lr)
dataloader = self.prepare_dataset(dataset_name, batch_size, image_size)
total_steps = len(dataloader) * num_epochs
with tqdm(total=total_steps, desc="Entraînement Lou") as pbar:
for epoch in range(num_epochs):
self.current_epoch = epoch
avg_loss = self.train_epoch(dataloader, pbar)
print(f"Epoch {epoch+1}/{num_epochs} - Loss: {avg_loss:.4f}")
# Sauvegarde périodique
if (epoch + 1) % save_every == 0 and repo_id and token:
self.save_checkpoint(repo_id, token, epoch + 1)
if progress_callback:
progress_callback(epoch + 1, num_epochs, avg_loss)
# Sauvegarde finale
if repo_id and token:
self.save_checkpoint(repo_id, token, "final")
return True
def save_checkpoint(self, repo_id, token, epoch):
"""Sauvegarde sur HF Hub"""
os.makedirs("checkpoints", exist_ok=True)
# Sauvegarde locale
checkpoint_path = f"checkpoints/lou_epoch_{epoch}.pt"
torch.save({
'epoch': epoch,
'model_state_dict': self.model.state_dict(),
'optimizer_state_dict': self.optimizer.state_dict() if self.optimizer else None,
}, checkpoint_path)
# Upload vers HF
try:
api = HfApi(token=token)
try:
create_repo(repo_id, private=False, token=token)
except:
pass
api.upload_file(
path_or_fileobj=checkpoint_path,
path_in_repo=f"lou_epoch_{epoch}.pt",
repo_id=repo_id,
token=token
)
print(f"✓ Checkpoint epoch {epoch} uploadé sur {repo_id}")
except Exception as e:
print(f"✗ Erreur upload: {e}")
def load_checkpoint(self, checkpoint_path):
"""Charge un checkpoint"""
checkpoint = torch.load(checkpoint_path, map_location=self.device)
self.model.load_state_dict(checkpoint['model_state_dict'])
if self.optimizer and checkpoint['optimizer_state_dict']:
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
self.current_epoch = checkpoint.get('epoch', 0)
print(f"✓ Checkpoint chargé (epoch {self.current_epoch})") |