DenozordecandClaude Sonnet 5.5
05537ac9e9
docker / build (push) Successful in 1m40s
- harvest_core: дубляж — тип, а не часть имени студии; split_multi_rows; audio_only_rows; kinozal сопоставляет все аудиодорожки - prepare_data: --max-dup, отсев остатков многотиповых строк - train: веса классов (корень из обратной частоты, потолок 4), --no-class-weights - исправлены регулярки с \b, записанные как символ BS (predict.py, audio18.py) - PARSE_VERSION 3: кэш разбора сбрасывается - docs/dataset.md, docs/api-reference.md обновлены Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
454 lines
23 KiB
Python
454 lines
23 KiB
Python
"""Обучение multi-task модели (тип озвучки + автор) и сравнение конфигураций.
|
||
|
||
Запуск:
|
||
python -m src.train # обе модели: 5-fold CV + финальные
|
||
python -m src.train --models tiny2 # только rubert-tiny2
|
||
python -m src.train --models base --folds 3 --epochs-override 5
|
||
|
||
Этапы для каждой модели:
|
||
1. 5-fold CV на train-части (группы: релиз + одинаковый текст) -> средние метрики;
|
||
2. финальное обучение на train (early stopping на 10% val-группах);
|
||
3. однократная оценка на held-out test -> reports/final_<имя>.json,
|
||
confusion matrix и кривые обучения -> reports/*.png;
|
||
4. чекпоинт -> models/<имя>/. Лучшая модель копируется в models/best.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import random
|
||
import re
|
||
import shutil
|
||
import time
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import torch
|
||
from sklearn.metrics import accuracy_score, confusion_matrix, f1_score
|
||
from torch.utils.data import DataLoader, Dataset
|
||
from transformers import AutoTokenizer, get_linear_schedule_with_warmup
|
||
|
||
from .model import DubTagger, predict_with_lang, save_model
|
||
from .spans import MAX_LEN, decode_span # noqa: F401 (прежнее место; сервису не нужны sklearn/matplotlib)
|
||
|
||
MODEL_CONFIGS = {
|
||
"tiny2": {"backbone": "cointegrated/rubert-tiny2", "lr": 1e-4,
|
||
"batch": 64, "epochs": 20},
|
||
"base": {"backbone": "ai-forever/ruBert-base", "lr": 3e-5,
|
||
"batch": 32, "epochs": 10},
|
||
# крупные профили (RTX 5070 Ti 16 ГБ: bf16 + batch 16 помещаются); в обучение по умолчанию не входят
|
||
"large": {"backbone": "ai-forever/ruBert-large", "lr": 2e-5,
|
||
"batch": 16, "epochs": 8},
|
||
"modern": {"backbone": "deepvk/RuModernBERT-base", "lr": 5e-5,
|
||
"batch": 32, "epochs": 8},
|
||
}
|
||
DEFAULT_MODELS = ["tiny2", "base"] # как в docs/training-windows.md: python -m src.train
|
||
PROFILE_INFO = {
|
||
"tiny2": "rubert-tiny2, 30M: самая быстрая, годится для ONNX/CPU",
|
||
"base": "ruBert-base, 178M: текущий рабочий выбор",
|
||
"large": "ruBert-large, 427M: та же семья, но втрое больше; ~3× дольше и тяжелее на CPU",
|
||
"modern": "RuModernBERT-base, 149M (2024): новая архитектура; ONNX-экспорт проекта её пока не поддерживает",
|
||
}
|
||
SEED = 42
|
||
|
||
|
||
def set_seed(seed: int) -> None:
|
||
random.seed(seed)
|
||
np.random.seed(seed)
|
||
torch.manual_seed(seed)
|
||
|
||
|
||
class Rows(Dataset):
|
||
def __init__(self, rows: list[dict]):
|
||
self.rows = rows
|
||
|
||
def __len__(self):
|
||
return len(self.rows)
|
||
|
||
def __getitem__(self, i):
|
||
return self.rows[i]
|
||
|
||
|
||
def make_collate(tokenizer, type2id: dict[str, int], max_len: int = MAX_LEN,
|
||
lang2id: dict[str, int] | None = None):
|
||
"""Токенизация батча + BIO-метки по символьным спанам (через offsets).
|
||
|
||
Возвращает (enc, types, bio, offsets): offsets — список списков (a, b)
|
||
по токенам, нужен для декодирования спана при оценке/инференсе.
|
||
"""
|
||
|
||
def collate(batch: list[dict]):
|
||
texts = [b["text"] for b in batch]
|
||
enc = tokenizer(texts, truncation=True, max_length=max_len,
|
||
padding=True, return_tensors="pt",
|
||
return_offsets_mapping=True)
|
||
offsets = enc.pop("offset_mapping").tolist()
|
||
# O по умолчанию (включая спецтокены — модель учится говорить «не автор»);
|
||
# -100 только для паддинга, он не входит в лосс
|
||
bio = torch.zeros_like(enc["input_ids"])
|
||
bio[enc["attention_mask"] == 0] = -100
|
||
for i, b in enumerate(batch):
|
||
span = b.get("span")
|
||
if not span:
|
||
continue
|
||
s, e = span
|
||
first = True
|
||
for t, (a, z) in enumerate(offsets[i]):
|
||
if z > a and a < e and z > s: # токен пересекает спан
|
||
bio[i, t] = 1 if first else 2 # B / I
|
||
first = False
|
||
types = torch.tensor([type2id[b["translation_type"]] for b in batch])
|
||
ages = torch.tensor([int(bool(b.get("is_18plus"))) for b in batch])
|
||
if lang2id: # неизвестный язык (None) = -100: в лосс и метрику не входит
|
||
langs = torch.tensor([lang2id.get(b.get("lang"), -100) if b.get("lang") else -100
|
||
for b in batch])
|
||
return enc, types, bio, ages, offsets, langs
|
||
return enc, types, bio, ages, offsets
|
||
|
||
return collate
|
||
|
||
|
||
def _NORM(s: str | None) -> str | None:
|
||
return re.sub(r"\s+", " ", s).strip().lower() if s else None
|
||
|
||
|
||
@torch.no_grad()
|
||
def evaluate(model, loader, device, id2type, per_row: list | None = None) -> dict:
|
||
"""Метрики на выборке: тип (acc, macro-F1), спан (exact), 18+ (acc, F1).
|
||
|
||
Exact match спана сравнивает нормализованный текст предсказанного спана
|
||
с текстом золотого спана (None == None считается совпадением).
|
||
joint_exact — строка распознана полностью: тип + автор + 18+.
|
||
"""
|
||
model.eval()
|
||
y_ids, p_ids = [], []
|
||
span_ok, y_age, p_age = [], [], []
|
||
y_lang, p_lang = [], []
|
||
for enc, types, _, ages, offsets, *rest in loader:
|
||
enc_dev = {k: v.to(device) for k, v in enc.items()}
|
||
with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device == "cuda"):
|
||
if model.lang_head is not None and rest:
|
||
type_logits, bio_logits, age_logits, lang_logits = predict_with_lang(model, **enc_dev)
|
||
y_lang += rest[0].tolist()
|
||
p_lang += lang_logits.argmax(-1).tolist()
|
||
else:
|
||
type_logits, bio_logits, age_logits = model.predict(**enc_dev)
|
||
p_ids += type_logits.argmax(-1).tolist()
|
||
y_ids += types.tolist()
|
||
p_age += age_logits.argmax(-1).tolist()
|
||
y_age += ages.tolist()
|
||
bio_pred = bio_logits.argmax(-1).cpu().numpy()
|
||
attn = enc["attention_mask"].tolist()
|
||
base = len(span_ok)
|
||
rows = loader.dataset.rows
|
||
for i in range(len(types)):
|
||
pred = decode_span(bio_pred[i], offsets[i], attn[i])
|
||
gold = rows[base + i].get("span")
|
||
pred_txt = rows[base + i]["text"][pred[0]:pred[1]] if pred else None
|
||
gold_txt = rows[base + i]["text"][gold[0]:gold[1]] if gold else None
|
||
span_ok.append(_NORM(pred_txt) == _NORM(gold_txt))
|
||
y = [id2type[v] for v in y_ids]
|
||
p = [id2type[v] for v in p_ids]
|
||
type_ok = [a == b for a, b in zip(y, p)]
|
||
age_ok = [a == b for a, b in zip(y_age, p_age)]
|
||
if per_row is not None: # 1 = строка распознана полностью (тип + автор + 18+), по порядку loader
|
||
per_row.extend(int(a and b and c) for a, b, c in zip(type_ok, span_ok, age_ok))
|
||
lang_pairs = [(a, b) for a, b in zip(y_lang, p_lang) if a != -100]
|
||
out_lang = ({"lang_accuracy": float(np.mean([a == b for a, b in lang_pairs])),
|
||
"lang_n": len(lang_pairs)} if lang_pairs else {})
|
||
return {
|
||
**out_lang,
|
||
"type_accuracy": float(np.mean(type_ok)) if type_ok else 0.0,
|
||
"type_macro_f1": f1_score(y, p, average="macro", zero_division=0),
|
||
"span_exact": float(np.mean(span_ok)) if span_ok else 0.0,
|
||
"age_accuracy": float(np.mean(age_ok)) if age_ok else 0.0,
|
||
"age_f1": f1_score(y_age, p_age, average="binary", zero_division=0)
|
||
if any(y_age) or any(p_age) else 1.0,
|
||
"joint_exact": float(np.mean([a and b and c for a, b, c in
|
||
zip(type_ok, span_ok, age_ok)]))
|
||
if type_ok else 0.0,
|
||
}
|
||
|
||
|
||
def class_weights(rows, type2id: dict[str, int], power: float = 0.5) -> torch.Tensor:
|
||
"""Веса классов типа: (N / (K · n_c)) ** power, нормированные так, что средний вес по примерам = 1.
|
||
|
||
power 0.5 (корень из обратной частоты) смягчает дисбаланс, не раздувая шум редчайших классов
|
||
(у «авт. многоголосый» несколько строк): полная обратная частота дала бы им вес в сотни раз больше.
|
||
Потолок веса — 4.
|
||
"""
|
||
counts = np.ones(len(type2id))
|
||
for r in rows:
|
||
counts[type2id[r["translation_type"]]] += 1
|
||
w = (counts.sum() / (len(counts) * counts)) ** power
|
||
w = w / (w * counts).sum() * counts.sum()
|
||
w = np.minimum(w, 4.0) # потолок: классы из единиц строк не должны перетягивать обучение
|
||
return torch.tensor(w, dtype=torch.float32)
|
||
|
||
|
||
def train_model(name: str, cfg: dict, train_rows, val_rows, test_rows,
|
||
type2id, device, epochs=None, seed=SEED, log_prefix="", lang2id=None,
|
||
use_class_weights: bool = True):
|
||
"""Обучение с early stopping на val; возвращает (модель, токенизатор, отчёт)."""
|
||
set_seed(seed)
|
||
tokenizer = AutoTokenizer.from_pretrained(cfg["backbone"])
|
||
if not tokenizer.is_fast: # спан автора считается по offsets: нужен быстрый токенизатор
|
||
raise RuntimeError(f"{cfg['backbone']}: нет быстрого токенизатора, offsets недоступны")
|
||
model = DubTagger(cfg["backbone"], n_types=len(type2id),
|
||
n_langs=len(lang2id or {})).to(device)
|
||
collate = make_collate(tokenizer, type2id, lang2id=lang2id)
|
||
train_loader = DataLoader(Rows(train_rows), batch_size=cfg["batch"],
|
||
shuffle=True, collate_fn=collate)
|
||
val_loader = DataLoader(Rows(val_rows), batch_size=128, collate_fn=collate)
|
||
test_loader = DataLoader(Rows(test_rows), batch_size=128, collate_fn=collate)
|
||
id2type = {i: t for t, i in type2id.items()}
|
||
if use_class_weights:
|
||
model.set_type_weights(class_weights(train_rows, type2id))
|
||
|
||
epochs = epochs or cfg["epochs"]
|
||
steps = len(train_loader) * epochs
|
||
optim = torch.optim.AdamW(model.parameters(), lr=cfg["lr"], weight_decay=0.01)
|
||
sched = get_linear_schedule_with_warmup(optim, int(steps * 0.1), steps)
|
||
|
||
best, best_state, history = -1.0, None, []
|
||
patience_left = 3
|
||
for epoch in range(epochs):
|
||
model.train()
|
||
losses = []
|
||
for enc, types, bio, ages, _, *rest in train_loader:
|
||
enc_dev = {k: v.to(device) for k, v in enc.items()}
|
||
with torch.autocast("cuda", dtype=torch.bfloat16,
|
||
enabled=device == "cuda"):
|
||
loss, _, _, _ = model(**enc_dev, type_labels=types.to(device),
|
||
bio_labels=bio.to(device),
|
||
age_labels=ages.to(device),
|
||
lang_labels=rest[0].to(device) if rest else None)
|
||
optim.zero_grad()
|
||
loss.backward()
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||
optim.step()
|
||
sched.step()
|
||
losses.append(loss.item())
|
||
m = evaluate(model, val_loader, device, id2type)
|
||
score = (m["type_macro_f1"] + m["span_exact"] + m["age_f1"]) / 3
|
||
history.append({"epoch": epoch, "loss": float(np.mean(losses)), **m})
|
||
print(f" {log_prefix}[{name}] epoch {epoch}: "
|
||
f"loss={np.mean(losses):.4f} type_f1={m['type_macro_f1']:.3f} "
|
||
f"span={m['span_exact']:.3f} age={m['age_accuracy']:.3f}"
|
||
+ (f" lang={m['lang_accuracy']:.3f} (n={m['lang_n']})" if "lang_accuracy" in m else ""))
|
||
if score > best:
|
||
best = score
|
||
best_state = {k: v.detach().cpu().clone()
|
||
for k, v in model.state_dict().items()}
|
||
patience_left = 3
|
||
else:
|
||
patience_left -= 1
|
||
if patience_left <= 0:
|
||
break
|
||
if best_state:
|
||
model.load_state_dict(best_state)
|
||
per_row: list = []
|
||
report = {"history": history,
|
||
"test": evaluate(model, test_loader, device, id2type, per_row=per_row),
|
||
"per_row": per_row}
|
||
return model, tokenizer, report
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser(description=__doc__,
|
||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||
ap.add_argument("--models", nargs="+", default=DEFAULT_MODELS,
|
||
choices=list(MODEL_CONFIGS))
|
||
ap.add_argument("--data", default="data/processed/dataset.json")
|
||
ap.add_argument("--label-maps", default=None, help="по умолчанию label_maps.json рядом с --data")
|
||
ap.add_argument("--folds", type=int, default=5, help="0 — без кросс-валидации (втрое-впятеро быстрее)")
|
||
ap.add_argument("--epochs-override", type=int, default=None)
|
||
ap.add_argument("--lr", type=float, default=None, help="заменяет lr профиля")
|
||
ap.add_argument("--batch", type=int, default=None, help="заменяет batch профиля")
|
||
ap.add_argument("--seed", type=int, default=SEED)
|
||
ap.add_argument("--no-class-weights", action="store_true",
|
||
help="потери типа без весов классов (по умолчанию: корень из обратной частоты)")
|
||
ap.add_argument("--runs", action="store_true",
|
||
help="каждая модель — отдельный запуск в --runs-dir; models/best, reports/final_* "
|
||
"и summary.md не трогаются (выбор лучшей — tools/collect.py, пункт 7)")
|
||
ap.add_argument("--runs-dir", default="models/runs")
|
||
ap.add_argument("--tag", default="", help="суффикс имени запуска")
|
||
args = ap.parse_args()
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
print(f"device: {device}" +
|
||
(f" ({torch.cuda.get_device_name(0)})" if device == "cuda" else ""))
|
||
|
||
rows = json.loads(Path(args.data).read_text(encoding="utf-8"))
|
||
maps_path = Path(args.label_maps) if args.label_maps else Path(args.data).with_name("label_maps.json")
|
||
label_maps = json.loads(maps_path.read_text(encoding="utf-8"))
|
||
type2id = {t: i for i, t in enumerate(label_maps["types"])}
|
||
lang2id = {l: i for i, l in enumerate(label_maps.get("langs") or [])} or None
|
||
id2type = {i: t for t, i in type2id.items()}
|
||
train_rows = [r for r in rows if r["split"] == "train"]
|
||
test_rows = [r for r in rows if r["split"] == "test"]
|
||
print(f"rows: train {len(train_rows)} / test {len(test_rows)}, "
|
||
f"classes {len(type2id)}")
|
||
|
||
from . import runs
|
||
fp = runs.fingerprint(test_rows)
|
||
Path("models").mkdir(exist_ok=True)
|
||
summary: dict[str, dict] = {}
|
||
|
||
for name in args.models:
|
||
cfg = dict(MODEL_CONFIGS[name])
|
||
if args.lr:
|
||
cfg["lr"] = args.lr
|
||
if args.batch:
|
||
cfg["batch"] = args.batch
|
||
t0 = time.time()
|
||
|
||
# --- 5-fold CV на train-части ---------------------------------
|
||
fold_ids = sorted({r["fold"] for r in train_rows})
|
||
cv_metrics = []
|
||
if args.folds:
|
||
print(f"\n=== {name} ({cfg['backbone']}): CV ===")
|
||
for f in fold_ids[:args.folds]:
|
||
tr = [r for r in train_rows if r["fold"] != f]
|
||
va = [r for r in train_rows if r["fold"] == f]
|
||
_, _, res = train_model(name, cfg, tr, va, va, type2id, device,
|
||
epochs=args.epochs_override,
|
||
seed=args.seed + f, log_prefix=f"cv{f} ", lang2id=lang2id,
|
||
use_class_weights=not args.no_class_weights)
|
||
cv_metrics.append(res["test"])
|
||
cv = ({k: float(np.mean([m[k] for m in cv_metrics])) for k in cv_metrics[0]}
|
||
if cv_metrics else None)
|
||
|
||
# --- финальная модель: train (90%) + val (10% групп) ------------
|
||
print(f"=== {name}: финальное обучение ===")
|
||
groups = sorted({r["group_id"] for r in train_rows})
|
||
rng = random.Random(args.seed)
|
||
rng.shuffle(groups)
|
||
val_groups = set(groups[:max(1, len(groups) // 10)])
|
||
tr = [r for r in train_rows if r["group_id"] not in val_groups]
|
||
va = [r for r in train_rows if r["group_id"] in val_groups]
|
||
model, tokenizer, res = train_model(
|
||
name, cfg, tr, va, test_rows, type2id, device,
|
||
epochs=args.epochs_override, seed=args.seed, lang2id=lang2id,
|
||
use_class_weights=not args.no_class_weights)
|
||
res["cv_mean"] = cv
|
||
res["train_seconds"] = round(time.time() - t0, 1)
|
||
per_row = res.pop("per_row")
|
||
summary[name] = res
|
||
|
||
if args.runs: # отдельный запуск: ничего существующего не перезаписывается
|
||
run_dir = Path(args.runs_dir) / f"{time.strftime('%Y%m%d_%H%M%S')}_{name}{args.tag}"
|
||
save_model(model, tokenizer, run_dir)
|
||
shutil.copy(maps_path, run_dir / "label_maps.json")
|
||
info = {"name": name, **cfg, "max_len": MAX_LEN, "seed": args.seed, "folds": args.folds,
|
||
"class_weights": not args.no_class_weights,
|
||
"created": time.strftime("%Y-%m-%d %H:%M:%S"), "fingerprint": fp,
|
||
"rows": {"train": len(train_rows), "test": len(test_rows)},
|
||
"final": {k: v for k, v in res.items() if k != "history"}}
|
||
(run_dir / "run.json").write_text(json.dumps(info, ensure_ascii=False, indent=1), encoding="utf-8")
|
||
(run_dir / "train_config.json").write_text(json.dumps(info, ensure_ascii=False, indent=1),
|
||
encoding="utf-8")
|
||
runs.store_eval(run_dir, fp, {"metrics": res["test"], "per_row": per_row, "n_unscorable": 0})
|
||
_plots(name, res, model, tokenizer, test_rows, type2id, id2type, device, lang2id, out=run_dir)
|
||
print(f"запуск сохранён: {run_dir}")
|
||
del model
|
||
if device == "cuda":
|
||
torch.cuda.empty_cache()
|
||
continue
|
||
|
||
model_dir = Path("models") / name
|
||
save_model(model, tokenizer, model_dir)
|
||
shutil.copy(maps_path, model_dir)
|
||
(model_dir / "train_config.json").write_text(json.dumps(
|
||
{"name": name, **cfg, "max_len": MAX_LEN, "seed": args.seed,
|
||
"class_weights": not args.no_class_weights,
|
||
"final": {k: v for k, v in res.items() if k != "history"}},
|
||
ensure_ascii=False, indent=1), encoding="utf-8")
|
||
runs.store_eval(model_dir, fp, {"metrics": res["test"], "per_row": per_row, "n_unscorable": 0})
|
||
Path(f"reports/final_{name}.json").write_text(
|
||
json.dumps(res, ensure_ascii=False, indent=1), encoding="utf-8")
|
||
|
||
_plots(name, res, model, tokenizer, test_rows, type2id, id2type, device, lang2id)
|
||
|
||
if args.runs:
|
||
return
|
||
|
||
# --- выбор лучшей и итоговый отчёт ---------------------------------
|
||
best_name = max(summary, key=lambda n: (summary[n]["test"]["joint_exact"],
|
||
summary[n]["test"]["type_macro_f1"]))
|
||
if Path("models/best").exists():
|
||
shutil.rmtree("models/best")
|
||
shutil.copytree(Path("models") / best_name, "models/best")
|
||
|
||
lines = ["# Сравнение моделей", "",
|
||
"| модель | type acc | type macro-F1 | span exact | age acc | joint exact | время, с |",
|
||
"|---|---|---|---|---|---|---|"]
|
||
for n, r in summary.items():
|
||
t = r["test"]
|
||
lines.append(f"| {n} | {t['type_accuracy']:.3f} | {t['type_macro_f1']:.3f} "
|
||
f"| {t['span_exact']:.3f} | {t['age_accuracy']:.3f} "
|
||
f"| {t['joint_exact']:.3f} | {r['train_seconds']} |")
|
||
lines += ["", f"Лучшая модель: **{best_name}** -> `models/best`"]
|
||
Path("reports/summary.md").write_text("\n".join(lines), encoding="utf-8")
|
||
print("\n" + "\n".join(lines))
|
||
|
||
|
||
def _plots(name, res, model, tokenizer, test_rows, type2id, id2type, device, lang2id=None, out="reports"):
|
||
"""Кривые обучения и confusion matrix на test (в каталог out)."""
|
||
out = Path(out)
|
||
out.mkdir(parents=True, exist_ok=True)
|
||
import matplotlib
|
||
matplotlib.use("Agg")
|
||
import matplotlib.pyplot as plt
|
||
|
||
hist = res["history"]
|
||
if hist:
|
||
fig, ax = plt.subplots(1, 2, figsize=(10, 4))
|
||
ax[0].plot([h["epoch"] for h in hist], [h["loss"] for h in hist])
|
||
ax[0].set_title("train loss")
|
||
ax[0].set_xlabel("epoch")
|
||
ax[1].plot([h["epoch"] for h in hist], [h["type_macro_f1"] for h in hist],
|
||
label="type macro-F1")
|
||
ax[1].plot([h["epoch"] for h in hist], [h["span_exact"] for h in hist],
|
||
label="span exact")
|
||
ax[1].set_xlabel("epoch")
|
||
ax[1].legend()
|
||
fig.tight_layout()
|
||
fig.savefig(out / f"history_{name}.png", dpi=120)
|
||
plt.close(fig)
|
||
|
||
loader = DataLoader(Rows(test_rows), batch_size=128,
|
||
collate_fn=make_collate(tokenizer, type2id, lang2id=lang2id))
|
||
y, p = [], []
|
||
with torch.no_grad():
|
||
for enc, types, _, _, _, *_ in loader:
|
||
enc_dev = {k: v.to(device) for k, v in enc.items()}
|
||
tl, _, _ = model.predict(**enc_dev)
|
||
y += types.tolist()
|
||
p += tl.argmax(-1).tolist()
|
||
y = [id2type[v] for v in y]
|
||
p = [id2type[v] for v in p]
|
||
labels = sorted(set(y) | set(p), key=lambda t: -y.count(t))
|
||
cm = confusion_matrix(y, p, labels=labels)
|
||
fig, ax = plt.subplots(figsize=(9, 7))
|
||
im = ax.imshow(cm, cmap="Blues")
|
||
ax.set_xticks(range(len(labels)), labels, rotation=45, ha="right", fontsize=8)
|
||
ax.set_yticks(range(len(labels)), labels, fontsize=8)
|
||
for i in range(len(labels)):
|
||
for j in range(len(labels)):
|
||
if cm[i, j]:
|
||
ax.text(j, i, cm[i, j], ha="center", va="center", fontsize=7,
|
||
color="white" if cm[i, j] > cm.max() / 2 else "black")
|
||
fig.colorbar(im)
|
||
ax.set_title(f"confusion matrix: {name} (test)")
|
||
ax.set_xlabel("предсказано")
|
||
ax.set_ylabel("истина")
|
||
fig.tight_layout()
|
||
fig.savefig(out / f"confusion_{name}.png", dpi=120)
|
||
plt.close(fig)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|