348 lines
13 KiB
Python
348 lines
13 KiB
Python
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,
|
|
} |