train.py
7,656 bytes
| 1 | import argparse |
|---|---|
| 2 | import csv |
| 3 | import random |
| 4 | from pathlib import Path |
| 5 | |
| 6 | import numpy as np |
| 7 | import pandas as pd |
| 8 | import timm |
| 9 | import torch |
| 10 | import torch.nn as nn |
| 11 | import yaml |
| 12 | from torch.utils.data import DataLoader, WeightedRandomSampler |
| 13 | |
| 14 | from .data import CountryDataset, build_transforms, scan_image_folder, stratified_split |
| 15 | from .model import CountryClassifier |
| 16 | |
| 17 | |
| 18 | def set_seed(seed): |
| 19 | random.seed(seed) |
| 20 | np.random.seed(seed) |
| 21 | torch.manual_seed(seed) |
| 22 | torch.cuda.manual_seed_all(seed) |
| 23 | |
| 24 | |
| 25 | def train_one_epoch(model, loader, criterion, optimizer, scaler, device, use_amp): |
| 26 | model.train() |
| 27 | loss_sum, correct, total = 0.0, 0, 0 |
| 28 | for images, targets in loader: |
| 29 | images = images.to(device, non_blocking=True) |
| 30 | targets = targets.to(device, non_blocking=True) |
| 31 | with torch.autocast(device_type=device.type, enabled=use_amp): |
| 32 | logits = model(images) |
| 33 | loss = criterion(logits, targets) |
| 34 | optimizer.zero_grad(set_to_none=True) |
| 35 | scaler.scale(loss).backward() |
| 36 | scaler.unscale_(optimizer) |
| 37 | torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| 38 | scaler.step(optimizer) |
| 39 | scaler.update() |
| 40 | loss_sum += loss.item() * targets.size(0) |
| 41 | correct += (logits.argmax(1) == targets).sum().item() |
| 42 | total += targets.size(0) |
| 43 | return loss_sum / total, correct / total |
| 44 | |
| 45 | |
| 46 | @torch.no_grad() |
| 47 | def evaluate(model, loader, criterion, device, use_amp): |
| 48 | model.eval() |
| 49 | loss_sum, top1, top5, total = 0.0, 0, 0, 0 |
| 50 | for images, targets in loader: |
| 51 | images = images.to(device, non_blocking=True) |
| 52 | targets = targets.to(device, non_blocking=True) |
| 53 | with torch.autocast(device_type=device.type, enabled=use_amp): |
| 54 | logits = model(images) |
| 55 | loss = criterion(logits, targets) |
| 56 | loss_sum += loss.item() * targets.size(0) |
| 57 | top1 += (logits.argmax(1) == targets).sum().item() |
| 58 | k = min(5, logits.size(1)) |
| 59 | top5 += (logits.topk(k, dim=1).indices == targets.unsqueeze(1)).any(1).sum().item() |
| 60 | total += targets.size(0) |
| 61 | return loss_sum / total, top1 / total, top5 / total |
| 62 | |
| 63 | |
| 64 | def main(): |
| 65 | ap = argparse.ArgumentParser(description="Treenib riigiklassifitseerija") |
| 66 | ap.add_argument("--config", default="configs/default.yaml") |
| 67 | args = ap.parse_args() |
| 68 | with open(args.config, encoding="utf-8") as f: |
| 69 | cfg = yaml.safe_load(f) |
| 70 | |
| 71 | d, m, t = cfg["data"], cfg["model"], cfg["train"] |
| 72 | set_seed(t["seed"]) |
| 73 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 74 | use_amp = device.type == "cuda" |
| 75 | out_dir = Path(t["out_dir"]) |
| 76 | out_dir.mkdir(parents=True, exist_ok=True) |
| 77 | |
| 78 | manifest = Path(d["manifest"]) |
| 79 | if manifest.exists(): |
| 80 | df = pd.read_csv(manifest) |
| 81 | print(f"Manifest: {len(df)} pilti") |
| 82 | else: |
| 83 | print(f"Manifesti {manifest} pole, skannin {d['raw_dir']} (soovitav on enne käivitada clean.py)") |
| 84 | df = scan_image_folder(d["raw_dir"]) |
| 85 | counts = df["country"].value_counts() |
| 86 | df = df[df["country"].isin(counts[counts >= d["min_per_class"]].index)].reset_index(drop=True) |
| 87 | |
| 88 | classes = sorted(df["country"].unique()) |
| 89 | class_to_idx = {c: i for i, c in enumerate(classes)} |
| 90 | print(f"{len(df)} pilti, {len(classes)} riiki, seade: {device}") |
| 91 | |
| 92 | splits = stratified_split(df, d["val_frac"], d["test_frac"], t["seed"]) |
| 93 | for name, part in splits.items(): |
| 94 | part.to_csv(out_dir / f"{name}.csv", index=False) |
| 95 | |
| 96 | model = CountryClassifier(len(classes), m["backbone"], m["dropout"]).to(device) |
| 97 | data_cfg = timm.data.resolve_model_data_config(model.backbone) |
| 98 | mean, std, image_size = data_cfg["mean"], data_cfg["std"], data_cfg["input_size"][-1] |
| 99 | |
| 100 | train_ds = CountryDataset(splits["train"], class_to_idx, build_transforms(mean, std, image_size, train=True)) |
| 101 | val_ds = CountryDataset(splits["val"], class_to_idx, build_transforms(mean, std, image_size, train=False)) |
| 102 | |
| 103 | # Kaalutud valim: iga riik jõuab batch'idesse võrdse tõenäosusega, |
| 104 | # muidu domineeriksid suurte piltide arvuga riigid. |
| 105 | train_counts = splits["train"]["country"].value_counts() |
| 106 | weights = splits["train"]["country"].map(lambda c: 1.0 / train_counts[c]).to_numpy() |
| 107 | sampler = WeightedRandomSampler( |
| 108 | torch.tensor(weights, dtype=torch.double), |
| 109 | num_samples=len(weights), |
| 110 | replacement=True, |
| 111 | generator=torch.Generator().manual_seed(t["seed"]), |
| 112 | ) |
| 113 | loader_kw = dict(batch_size=t["batch_size"], num_workers=d["num_workers"], |
| 114 | pin_memory=use_amp, persistent_workers=d["num_workers"] > 0) |
| 115 | train_loader = DataLoader(train_ds, sampler=sampler, drop_last=True, **loader_kw) |
| 116 | val_loader = DataLoader(val_ds, shuffle=False, **loader_kw) |
| 117 | |
| 118 | criterion = nn.CrossEntropyLoss(label_smoothing=t["label_smoothing"]) |
| 119 | scaler = torch.amp.GradScaler(device.type, enabled=use_amp) |
| 120 | |
| 121 | log_path = out_dir / "log.csv" |
| 122 | with open(log_path, "w", newline="", encoding="utf-8") as f: |
| 123 | csv.writer(f).writerow( |
| 124 | ["phase", "epoch", "lr", "train_loss", "train_acc", "val_loss", "val_top1", "val_top5"] |
| 125 | ) |
| 126 | |
| 127 | best_top1 = 0.0 |
| 128 | |
| 129 | def log_and_save(phase, epoch, optimizer, train_loss, train_acc): |
| 130 | nonlocal best_top1 |
| 131 | val_loss, val_top1, val_top5 = evaluate(model, val_loader, criterion, device, use_amp) |
| 132 | lr = optimizer.param_groups[0]["lr"] |
| 133 | with open(log_path, "a", newline="", encoding="utf-8") as f: |
| 134 | csv.writer(f).writerow( |
| 135 | [phase, epoch, f"{lr:.2e}", f"{train_loss:.4f}", f"{train_acc:.4f}", |
| 136 | f"{val_loss:.4f}", f"{val_top1:.4f}", f"{val_top5:.4f}"] |
| 137 | ) |
| 138 | print(f"[{phase}] epohh {epoch}: train_acc={train_acc:.3f} " |
| 139 | f"val_top1={val_top1:.3f} val_top5={val_top5:.3f}") |
| 140 | if val_top1 > best_top1: |
| 141 | best_top1 = val_top1 |
| 142 | torch.save( |
| 143 | {"model": model.state_dict(), "classes": classes, |
| 144 | "backbone": m["backbone"], "dropout": m["dropout"], "val_top1": val_top1}, |
| 145 | out_dir / "best.pt", |
| 146 | ) |
| 147 | |
| 148 | # Faas 1: backbone külmutatud, treenime ainult klassifitseerimispead. |
| 149 | # Suvaliselt initsialiseeritud pea gradiendid lõhuksid eeltreenitud kaale. |
| 150 | model.freeze_backbone() |
| 151 | optimizer = torch.optim.AdamW(model.head.parameters(), lr=t["lr_head"], |
| 152 | weight_decay=t["weight_decay"]) |
| 153 | for epoch in range(1, t["epochs_head"] + 1): |
| 154 | train_loss, train_acc = train_one_epoch(model, train_loader, criterion, |
| 155 | optimizer, scaler, device, use_amp) |
| 156 | log_and_save("head", epoch, optimizer, train_loss, train_acc) |
| 157 | |
| 158 | # Faas 2: viimased transformeri plokid lahti, madal LR ja koosinusgraafik. |
| 159 | model.set_finetune_mode(t["unfreeze_blocks"]) |
| 160 | backbone_params = [p for p in model.backbone.parameters() if p.requires_grad] |
| 161 | optimizer = torch.optim.AdamW( |
| 162 | [{"params": backbone_params, "lr": t["lr_backbone"]}, |
| 163 | {"params": model.head.parameters(), "lr": t["lr_head_finetune"]}], |
| 164 | weight_decay=t["weight_decay"], |
| 165 | ) |
| 166 | scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=t["epochs_finetune"]) |
| 167 | for epoch in range(1, t["epochs_finetune"] + 1): |
| 168 | train_loss, train_acc = train_one_epoch(model, train_loader, criterion, |
| 169 | optimizer, scaler, device, use_amp) |
| 170 | log_and_save("finetune", epoch, optimizer, train_loss, train_acc) |
| 171 | scheduler.step() |
| 172 | |
| 173 | print(f"Valmis. Parim val top-1: {best_top1:.3f}, checkpoint: {out_dir / 'best.pt'}") |
| 174 | |
| 175 | |
| 176 | if __name__ == "__main__": |
| 177 | main() |
| 178 | |