| """Trigger boundary 학습의 class-balance와 negative-sampling 계약. |
| |
| 이 모듈은 Trigger neural Head를 소유하지 않는다. 동일 logits/target에 적용할 학습 목적만 |
| 교체하며, production 기본값은 full-source all-negative BCE와 ``pos_weight=4``다. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from ..contracts import TriggerTrainingConfig |
|
|
|
|
| class TriggerBoundaryLossStrategy(nn.Module): |
| """동일 boundary logits에 weighted BCE와 deterministic negative sampling을 적용한다.""" |
|
|
| def __init__(self, config: TriggerTrainingConfig | None = None) -> None: |
| super().__init__() |
| self.config = config or TriggerTrainingConfig() |
|
|
| def forward( |
| self, |
| start_logits: torch.Tensor, |
| end_logits: torch.Tensor, |
| start_target: torch.Tensor, |
| end_target: torch.Tensor, |
| token_mask: torch.BoolTensor, |
| *, |
| sample_step: int = 0, |
| ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: |
| if ( |
| start_logits.shape != end_logits.shape |
| or start_target.shape != start_logits.shape |
| or end_target.shape != end_logits.shape |
| or token_mask.shape != start_logits.shape[:-1] |
| ): |
| raise ValueError("Trigger boundary loss tensor shapes differ") |
| if sample_step < 0: |
| raise ValueError("Trigger sample_step must be non-negative") |
|
|
| valid = token_mask.unsqueeze(-1).expand_as(start_logits) |
| start_loss, start_stats = self._one_boundary( |
| start_logits, |
| start_target, |
| valid, |
| sample_step=sample_step, |
| stream=0, |
| ) |
| end_loss, end_stats = self._one_boundary( |
| end_logits, |
| end_target, |
| valid, |
| sample_step=sample_step, |
| stream=1, |
| ) |
| diagnostics = { |
| "trigger/start_loss": start_loss, |
| "trigger/end_loss": end_loss, |
| "trigger/effective_pos_weight": start_logits.new_tensor( |
| self.config.pos_weight |
| ), |
| } |
| diagnostics.update( |
| {f"trigger/start_{name}": value for name, value in start_stats.items()} |
| ) |
| diagnostics.update( |
| {f"trigger/end_{name}": value for name, value in end_stats.items()} |
| ) |
| diagnostics["trigger/positive_boundary_count"] = ( |
| start_stats["positive_count"] + end_stats["positive_count"] |
| ) |
| diagnostics["trigger/negative_boundary_count"] = ( |
| start_stats["negative_count"] + end_stats["negative_count"] |
| ) |
| diagnostics["trigger/sampled_negative_count"] = ( |
| start_stats["sampled_negative_count"] |
| + end_stats["sampled_negative_count"] |
| ) |
| return start_loss + end_loss, diagnostics |
|
|
| def _one_boundary( |
| self, |
| logits: torch.Tensor, |
| target: torch.Tensor, |
| valid: torch.BoolTensor, |
| *, |
| sample_step: int, |
| stream: int, |
| ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: |
| positive = valid & target.gt(0.5) |
| negative = valid & ~positive |
| selected_negative = self._sample_negative( |
| negative, |
| positive_count=int(positive.sum().item()), |
| sample_step=sample_step, |
| stream=stream, |
| ) |
| selected = positive | selected_negative |
| if not torch.any(selected): |
| loss = logits.sum() * 0.0 |
| else: |
| arguments: dict[str, object] = {} |
| if self.config.pos_weight != 1.0: |
| arguments["pos_weight"] = logits.new_tensor(self.config.pos_weight) |
| loss = F.binary_cross_entropy_with_logits( |
| logits[selected], |
| target[selected].to(logits.dtype), |
| **arguments, |
| ) |
| return loss, { |
| "positive_count": logits.new_tensor(float(positive.sum().item())), |
| "negative_count": logits.new_tensor(float(negative.sum().item())), |
| "sampled_negative_count": logits.new_tensor( |
| float(selected_negative.sum().item()) |
| ), |
| } |
|
|
| def _sample_negative( |
| self, |
| negative: torch.BoolTensor, |
| *, |
| positive_count: int, |
| sample_step: int, |
| stream: int, |
| ) -> torch.BoolTensor: |
| ratio = self.config.negative_to_positive_ratio |
| if ratio is None: |
| return negative |
| negative_flat = torch.nonzero( |
| negative.detach().to(device="cpu").reshape(-1), as_tuple=False |
| ).flatten() |
| |
| requested = math.ceil(ratio * max(positive_count, 1)) |
| count = min(int(negative_flat.numel()), requested) |
| selected_cpu = torch.zeros(negative.numel(), dtype=torch.bool) |
| if count: |
| generator = torch.Generator(device="cpu") |
| mixed_seed = ( |
| self.config.sampling_seed |
| + 1_000_003 * sample_step |
| + 97_409 * stream |
| ) % (2**63 - 1) |
| generator.manual_seed(mixed_seed) |
| order = torch.randperm( |
| negative_flat.numel(), generator=generator |
| )[:count] |
| selected_cpu[negative_flat[order]] = True |
| return selected_cpu.reshape(negative.shape).to(device=negative.device) |
|
|
