Initial commit
This commit is contained in:
348
train/trainer.py
Normal file
348
train/trainer.py
Normal file
@@ -0,0 +1,348 @@
|
||||
import random
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from PIL import Image
|
||||
from sklearn.model_selection import train_test_split
|
||||
from sklearn.metrics import classification_report
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from torch.optim import AdamW
|
||||
from transformers import (
|
||||
LayoutLMv3Processor,
|
||||
LayoutLMv3ForSequenceClassification,
|
||||
get_scheduler,
|
||||
)
|
||||
from tqdm import tqdm
|
||||
from loguru import logger
|
||||
|
||||
import config
|
||||
from train.schemas import TrainRequest, TrainStatus
|
||||
from utils.cache import (
|
||||
load_cache, save_label2id,
|
||||
load_replay_buffer, update_replay_buffer, load_label2id
|
||||
)
|
||||
|
||||
# 학습 상태 전역
|
||||
train_status = TrainStatus(
|
||||
running=False, epoch=0, total=0, loss=0.0, message="대기 중"
|
||||
)
|
||||
stop_requested = False
|
||||
|
||||
|
||||
# ── Dataset ───────────────────────────────────────
|
||||
class InvoiceDataset(Dataset):
|
||||
def __init__(self, samples: list, processor):
|
||||
self.samples = samples
|
||||
self.processor = processor
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
s = self.samples[idx]
|
||||
image = Image.open(s["image_path"]).convert("RGB")
|
||||
encoding = self.processor(
|
||||
image, text=s["words"], boxes=s["boxes"],
|
||||
return_tensors="pt", truncation=True,
|
||||
padding="max_length", max_length=config.MAX_LEN,
|
||||
)
|
||||
return {
|
||||
"input_ids": encoding["input_ids"].squeeze(),
|
||||
"attention_mask": encoding["attention_mask"].squeeze(),
|
||||
"bbox": encoding["bbox"].squeeze(),
|
||||
"pixel_values": encoding["pixel_values"].squeeze(),
|
||||
"labels": torch.tensor(s["label"], dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
# ── Fisher (EWC) ──────────────────────────────────
|
||||
def compute_fisher(model, loader) -> dict:
|
||||
fisher = {n: torch.zeros_like(p)
|
||||
for n, p in model.named_parameters() if p.requires_grad}
|
||||
model.eval()
|
||||
for batch in loader:
|
||||
model.zero_grad()
|
||||
out = model(
|
||||
input_ids = batch["input_ids"].to(config.DEVICE),
|
||||
attention_mask = batch["attention_mask"].to(config.DEVICE),
|
||||
bbox = batch["bbox"].to(config.DEVICE),
|
||||
pixel_values = batch["pixel_values"].to(config.DEVICE),
|
||||
labels = batch["labels"].to(config.DEVICE),
|
||||
)
|
||||
out.loss.backward()
|
||||
for n, p in model.named_parameters():
|
||||
if p.requires_grad and p.grad is not None:
|
||||
fisher[n] += p.grad.pow(2)
|
||||
return {n: f / max(len(loader), 1) for n, f in fisher.items()}
|
||||
|
||||
|
||||
# ── 검증 ──────────────────────────────────────────
|
||||
def evaluate_loader(model, loader) -> float:
|
||||
model.eval()
|
||||
correct, total = 0, 0
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
out = model(
|
||||
input_ids = batch["input_ids"].to(config.DEVICE),
|
||||
attention_mask = batch["attention_mask"].to(config.DEVICE),
|
||||
bbox = batch["bbox"].to(config.DEVICE),
|
||||
pixel_values = batch["pixel_values"].to(config.DEVICE),
|
||||
)
|
||||
preds = out.logits.argmax(dim=-1).cpu()
|
||||
correct += (preds == batch["labels"]).sum().item()
|
||||
total += len(batch["labels"])
|
||||
return correct / total if total else 0.0
|
||||
|
||||
|
||||
def _get_latest_model_path():
|
||||
"""
|
||||
날짜/차수 폴더에서 가장 최신 모델 경로 반환
|
||||
files/models/20260305/0002/ ← 최신
|
||||
"""
|
||||
base = config.MODEL_BASE_PATH
|
||||
all_models = sorted([
|
||||
d for d in base.rglob("config.json")
|
||||
])
|
||||
if not all_models:
|
||||
return None
|
||||
|
||||
print(all_models[-1].parent)
|
||||
return all_models[-1].parent
|
||||
|
||||
|
||||
# ── 학습 루프 ─────────────────────────────────────
|
||||
def run_train_loop(model, processor, train_s, val_s, req: TrainRequest,
|
||||
label2id: dict, strategy: str, save_path: Path):
|
||||
global train_status
|
||||
|
||||
train_loader = DataLoader(
|
||||
InvoiceDataset(train_s, processor),
|
||||
batch_size=req.batch_size, shuffle=True
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
InvoiceDataset(val_s, processor),
|
||||
batch_size=req.batch_size
|
||||
)
|
||||
optimizer = AdamW(model.parameters(), lr=req.lr)
|
||||
scheduler = get_scheduler(
|
||||
"cosine", optimizer=optimizer,
|
||||
num_warmup_steps=len(train_loader),
|
||||
num_training_steps=len(train_loader) * req.epochs,
|
||||
)
|
||||
|
||||
# EWC 사전 계산
|
||||
ewc_params, ewc_fisher = {}, {}
|
||||
if strategy == "ewc":
|
||||
ewc_params = {n: p.clone().detach()
|
||||
for n, p in model.named_parameters() if p.requires_grad}
|
||||
ewc_fisher = compute_fisher(model, val_loader)
|
||||
logger.info("EWC Fisher 계산 완료")
|
||||
|
||||
best_acc = 0.0
|
||||
total_batches = len(train_loader) * req.epochs
|
||||
|
||||
for epoch in range(req.epochs):
|
||||
model.train()
|
||||
total_loss = 0
|
||||
|
||||
pbar = tqdm(
|
||||
train_loader,
|
||||
desc=f"Epoch {epoch+1}/{req.epochs}",
|
||||
ncols=100,
|
||||
leave=True,
|
||||
)
|
||||
for batch_idx, batch in enumerate(pbar):
|
||||
if stop_requested:
|
||||
logger.info("학습 중단 요청 감지")
|
||||
return
|
||||
|
||||
out = model(
|
||||
input_ids = batch["input_ids"].to(config.DEVICE),
|
||||
attention_mask = batch["attention_mask"].to(config.DEVICE),
|
||||
bbox = batch["bbox"].to(config.DEVICE),
|
||||
pixel_values = batch["pixel_values"].to(config.DEVICE),
|
||||
labels = batch["labels"].to(config.DEVICE),
|
||||
)
|
||||
loss = out.loss
|
||||
if strategy == "ewc" and ewc_fisher:
|
||||
for n, p in model.named_parameters():
|
||||
if n in ewc_fisher:
|
||||
loss += req.ewc_lambda * (
|
||||
ewc_fisher[n] * (p - ewc_params[n]).pow(2)
|
||||
).sum()
|
||||
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
total_loss += loss.item()
|
||||
|
||||
pbar.set_postfix(loss=f"{loss.item():.4f}")
|
||||
|
||||
# 배치 단위 상태 업데이트
|
||||
done_batches = epoch * len(train_loader) + (batch_idx + 1)
|
||||
train_status.epoch = epoch + 1
|
||||
train_status.total = req.epochs
|
||||
train_status.loss = round(loss.item(), 4)
|
||||
train_status.batch = done_batches # 완료된 전체 배치
|
||||
train_status.total_batches = total_batches # 전체 배치 수
|
||||
train_status.message = (
|
||||
f"Epoch {epoch+1}/{req.epochs} | "
|
||||
f"batch {batch_idx+1}/{len(train_loader)} | "
|
||||
f"loss={loss.item():.4f}"
|
||||
)
|
||||
|
||||
avg_loss = total_loss / len(train_loader)
|
||||
val_acc = evaluate_loader(model, val_loader)
|
||||
|
||||
msg = (
|
||||
f"Epoch {epoch+1}/{req.epochs} | "
|
||||
f"loss={avg_loss:.4f} | val_acc={val_acc:.4f}"
|
||||
)
|
||||
tqdm.write(f"[결과] {msg}")
|
||||
logger.info(msg)
|
||||
|
||||
train_status.epoch = epoch + 1
|
||||
train_status.loss = round(avg_loss, 4)
|
||||
train_status.val_acc = round(val_acc, 4)
|
||||
train_status.message = msg
|
||||
|
||||
if val_acc > best_acc:
|
||||
best_acc = val_acc
|
||||
model.save_pretrained(str(save_path))
|
||||
processor.save_pretrained(str(save_path))
|
||||
save_label2id(label2id)
|
||||
logger.info(f"모델 저장 (val_acc={val_acc:.4f}) → {save_path}")
|
||||
|
||||
|
||||
# ── 학습 진입점 ───────────────────────────────────
|
||||
def run_train(req: TrainRequest):
|
||||
global train_status
|
||||
stop_requested = False
|
||||
train_status = TrainStatus(
|
||||
running=True, epoch=0, total=req.epochs,
|
||||
loss=0.0, message="학습 준비 중"
|
||||
)
|
||||
try:
|
||||
from utils.cache import load_label2id
|
||||
|
||||
# 날짜/차수 폴더 생성
|
||||
save_path = config.get_model_save_path()
|
||||
logger.info(f"모델 저장 경로: {save_path}")
|
||||
train_status.message = f"저장 경로: {save_path}"
|
||||
|
||||
samples = load_cache()
|
||||
label2id = load_label2id()
|
||||
id2label = {v: k for k, v in label2id.items()}
|
||||
|
||||
processor_path = (
|
||||
str(config.MODEL_BASE_PATH / "preprocessor_config.json")
|
||||
if (config.MODEL_BASE_PATH / "preprocessor_config.json").exists()
|
||||
else config.MODEL_NAME
|
||||
)
|
||||
# 최신 모델 폴더 찾기 (이어 학습)
|
||||
latest_model = _get_latest_model_path()
|
||||
model_path = str(latest_model) if latest_model else config.MODEL_NAME
|
||||
logger.info(f"기반 모델: {model_path}")
|
||||
|
||||
processor = LayoutLMv3Processor.from_pretrained(
|
||||
latest_model and str(latest_model) or config.MODEL_NAME,
|
||||
apply_ocr=False
|
||||
)
|
||||
model = LayoutLMv3ForSequenceClassification.from_pretrained(
|
||||
model_path,
|
||||
num_labels=len(label2id),
|
||||
id2label=id2label,
|
||||
label2id=label2id,
|
||||
ignore_mismatched_sizes=True,
|
||||
).to(config.DEVICE)
|
||||
|
||||
# 전략 자동 선택
|
||||
strategy = req.strategy
|
||||
if strategy == "auto":
|
||||
buf = load_replay_buffer()
|
||||
existing_types = set(buf.keys())
|
||||
new_types = {id2label[s["label"]] for s in samples}
|
||||
strategy = "replay" if (new_types - existing_types) else "ewc"
|
||||
train_status.message = f"전략: {strategy}"
|
||||
logger.info(f"학습 전략: {strategy}")
|
||||
|
||||
train_s, val_s = train_test_split(
|
||||
samples, test_size=0.2, random_state=42,
|
||||
stratify=[s["label"] for s in samples],
|
||||
)
|
||||
|
||||
if strategy == "replay":
|
||||
replay_s = [s for v in load_replay_buffer().values() for s in v]
|
||||
combined = train_s + replay_s
|
||||
random.shuffle(combined)
|
||||
run_train_loop(model, processor, combined, val_s, req, label2id, strategy, save_path)
|
||||
update_replay_buffer(train_s, id2label)
|
||||
else:
|
||||
run_train_loop(model, processor, train_s, val_s, req, label2id, strategy, save_path)
|
||||
update_replay_buffer(train_s, id2label)
|
||||
|
||||
train_status.message = f"학습 완료 → {save_path}"
|
||||
logger.info(f"학습 완료: {save_path}")
|
||||
|
||||
except Exception:
|
||||
train_status.message = f"오류: {traceback.format_exc()}"
|
||||
logger.exception("학습 중 예외 발생")
|
||||
finally:
|
||||
stop_requested = False
|
||||
train_status.running = False
|
||||
|
||||
|
||||
# ── 평가 진입점 ───────────────────────────────────
|
||||
def run_evaluate(test_size: float = 0.2) -> dict:
|
||||
samples = load_cache()
|
||||
label2id = load_label2id()
|
||||
id2label = {v: k for k, v in label2id.items()}
|
||||
|
||||
_, test_s = train_test_split(
|
||||
samples, test_size=test_size, random_state=42,
|
||||
stratify=[s["label"] for s in samples],
|
||||
)
|
||||
|
||||
model_path = _get_latest_model_path()
|
||||
if not model_path:
|
||||
raise RuntimeError("저장된 모델이 없습니다.")
|
||||
|
||||
processor = LayoutLMv3Processor.from_pretrained(
|
||||
str(model_path), apply_ocr=False
|
||||
)
|
||||
model = LayoutLMv3ForSequenceClassification.from_pretrained(
|
||||
str(model_path)
|
||||
).to(config.DEVICE)
|
||||
model.eval()
|
||||
|
||||
loader = DataLoader(InvoiceDataset(test_s, processor), batch_size=config.BATCH_SIZE)
|
||||
all_preds, all_labels = [], []
|
||||
with torch.no_grad():
|
||||
for batch in loader:
|
||||
out = model(
|
||||
input_ids = batch["input_ids"].to(config.DEVICE),
|
||||
attention_mask = batch["attention_mask"].to(config.DEVICE),
|
||||
bbox = batch["bbox"].to(config.DEVICE),
|
||||
pixel_values = batch["pixel_values"].to(config.DEVICE),
|
||||
)
|
||||
all_preds.extend(out.logits.argmax(dim=-1).cpu().tolist())
|
||||
all_labels.extend(batch["labels"].tolist())
|
||||
|
||||
target_names = [id2label[i] for i in range(len(id2label))]
|
||||
report_dict = classification_report(
|
||||
all_labels, all_preds, target_names=target_names, output_dict=True
|
||||
)
|
||||
report_str = classification_report(
|
||||
all_labels, all_preds, target_names=target_names
|
||||
)
|
||||
logger.info(f"평가 완료 | accuracy={report_dict['accuracy']:.4f}")
|
||||
return {
|
||||
"test_samples": len(test_s),
|
||||
"accuracy": round(report_dict["accuracy"], 4),
|
||||
"report": report_dict,
|
||||
"report_text": report_str,
|
||||
}
|
||||
Reference in New Issue
Block a user