# Pitfalls

Common mistakes when using `zyplib.train` and how to avoid them.

## #1 Missing metric recorders

**Symptom**: `TypeError` or `AttributeError` crash inside `train_epoch` at the first batch.

**Cause**: `train_epoch` calls `_require_train_metric_recorder(ctx)` which casts to `TorchMetricRecorder` unconditionally. If `ctx.train_metric_recorder` is None, it crashes.

**Wrong**:
```python
ctx = FabricTrainingContext(model=model, optimizer=optimizer, fabric=fabric, fn_loss=loss_fn)
# train_metric_recorder and val_metric_recorder are both None!
procedure.train_loop(ctx, train_loader, val_loader=val_loader)  # CRASH
```

**Correct**:
```python
ctx = FabricTrainingContext(
    model=model, optimizer=optimizer, fabric=fabric, fn_loss=loss_fn,
    train_metric_recorder=TorchMetricRecorder(clf_metrics_builder(), before_metric='sigmoid'),
    val_metric_recorder=TorchMetricRecorder(clf_metrics_builder(), before_metric='sigmoid'),
)
```

**Rule**: Always provide both `train_metric_recorder` and `val_metric_recorder` when using `train_loop` or `BasicFabricTrainer`.

## #2 Using fabric.launch() instead of ctx.launch()

**Symptom**: `RuntimeError: Expected all tensors to be on the same device` during metric computation.

**Cause**: `ctx.launch()` does two things: (1) calls `fabric.launch()`, (2) moves metric recorders to `fabric.device`. If you only call `fabric.launch()`, torchmetrics stay on CPU while training tensors are on GPU.

**Wrong**:
```python
fabric.launch()  # metrics stay on CPU
```

**Correct**:
```python
ctx.launch()  # moves metrics to correct device
```

**Exception**: `BasicFabricTrainer` with `launch_fabric_internally=True` (default) handles this automatically.

## #3 Wrong before_metric for task type

**Symptom**: All metric values (accuracy, F1, etc.) are nonsensical — near 0 or near random.

**Cause**: `before_metric` must match the loss function's output expectations.

| Loss function | Required before_metric | Why |
|--------------|----------------------|-----|
| `BCEWithLogitsLoss` | `'sigmoid'` | Model outputs raw logits, metrics need probabilities |
| `CrossEntropyLoss` | `'softmax'` | Model outputs raw logits, metrics need probabilities |
| `BCELoss` | `None` | Model already outputs probabilities |
| Custom (EDL) | depends | EDL returns logits, needs activation |

**Wrong**:
```python
fn_loss = nn.BCEWithLogitsLoss()
recorder = TorchMetricRecorder(clf_metrics_builder(), before_metric=None)  # BAD
```

**Correct**:
```python
fn_loss = nn.BCEWithLogitsLoss()
recorder = TorchMetricRecorder(clf_metrics_builder(), before_metric='sigmoid')
```

## #4 Double model/optimizer setup

**Symptom**: `RuntimeError` about already-wrapped modules, or silent performance degradation.

**Cause**: `FabricTrainingContext.__post_init__` automatically calls `fabric.setup(model)` and `fabric.setup_optimizers(optimizer)` when `fabric` is provided. Calling these manually before creating the context causes double-wrapping.

**Wrong**:
```python
model = fabric.setup(model)          # first wrap
optimizer = fabric.setup_optimizers(optimizer)  # first wrap
ctx = FabricTrainingContext(model=model, optimizer=optimizer, fabric=fabric)  # wraps AGAIN
```

**Correct**:
```python
# Let FabricTrainingContext handle setup
ctx = FabricTrainingContext(model=model, optimizer=optimizer, fabric=fabric)
```

Source: `_base.py:117-153` — `__post_init__` checks `isinstance(_FabricModule)` before wrapping, but passing already-wrapped objects can still cause issues with optimizer state.

## #5 fn_loss resolution order

**Symptom**: `ValueError: '损失函数未指定'` even though you think you set it.

**Cause**: The loss function resolution has a specific priority order:
1. `ctx.fn_loss` (set on context)
2. `fn_loss` parameter passed to `train_loop` (deprecated, backward compat)
3. If both are None and no custom `fn_train_step` is provided → error

**Wrong**:
```python
ctx = FabricTrainingContext(model=model, optimizer=optimizer, fabric=fabric)
# fn_loss not set anywhere!
procedure.train_loop(ctx, train_loader)  # ERROR
```

**Correct** (preferred):
```python
ctx = FabricTrainingContext(model=model, optimizer=optimizer, fabric=fabric, fn_loss=nn.BCEWithLogitsLoss())
```

**Also correct** (deprecated):
```python
procedure.train_loop(ctx, train_loader, fn_loss=nn.BCEWithLogitsLoss())
```

**Note**: If you provide a custom `fn_train_step` that handles loss internally (like EDL/GAN), `fn_loss` can be None — but you'll get a warning unless the custom step is provided.

## #6 scalar_metrics name mismatch

**Symptom**: Loss values appear in console (from step function) but are missing from epoch-level logs and TensorBoard.

**Cause**: `TorchMetricRecorder.scalar_metrics` only tracks keys it knows about. Default is `'loss'`. If your step returns different keys, they're silently dropped (with a one-time warning).

**Wrong**:
```python
def my_step(ctx, batch, **kwargs):
    ...
    return y_pred, {'total_loss': loss.item(), 'reg_loss': reg.item()}

recorder = TorchMetricRecorder(metrics, scalar_metrics='loss')  # only tracks 'loss'
# 'total_loss' and 'reg_loss' are IGNORED
```

**Correct**:
```python
recorder = TorchMetricRecorder(metrics, scalar_metrics=['total_loss', 'reg_loss'])
# Or use auto mode:
recorder = TorchMetricRecorder(metrics, scalar_metrics='auto')
```

## #7 Forgetting val_metric_recorder for validation/test

**Symptom**: Crash during validation or `procedure.test()`.

**Cause**: `val_epoch` and `procedure.test` both use `_require_val_metric_recorder(ctx)`. If `val_metric_recorder` is None, it crashes.

**Wrong**:
```python
ctx = FabricTrainingContext(
    ...,
    train_metric_recorder=recorder,
    # val_metric_recorder not set!
)
procedure.train_loop(ctx, train_loader, val_loader=val_loader)  # CRASH on first val epoch
```

**Correct**:
```python
ctx = FabricTrainingContext(
    ...,
    train_metric_recorder=new_recorder(),
    val_metric_recorder=new_recorder(),  # separate instance!
)
```

**Note**: Train and val recorders should be **separate instances** (not the same object) because `compute()` resets internal state.

## #8 LR scheduler not registered as callback

**Symptom**: Learning rate never changes during training.

**Cause**: Simply creating a scheduler doesn't integrate it with the training loop. It must be registered as a `LRSchedulerCallback`.

**Wrong**:
```python
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30)
# scheduler exists but is never stepped!
procedure.train_loop(ctx, train_loader)
```

**Correct**:
```python
from zyplib.train.common.callbacks import LRSchedulerCallback
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30)
fabric = default_fabric(
    fn_custom_callback=lambda v: [LRSchedulerCallback(scheduler)]
)
```

**Backward compat** (deprecated): passing `lr_scheduler=scheduler` to `train_loop` auto-injects the callback, but this is deprecated.

## #9 Checkpoint monitor metric name mismatch

**Symptom**: Checkpoints are never saved (or saved every epoch regardless of performance).

**Cause**: `SaveCheckpoints.monitor` must match an actual metric name in the epoch metrics dict. If the name doesn't match, `current_score` is None and saving is skipped.

**Wrong**:
```python
fabric = default_fabric(monitor='accuracy')  # but metric is named 'acc' in MetricCollection
```

**Correct**:
```python
# Check what clf_metrics_builder actually names the metrics:
# It uses 'accuracy', 'precision', 'recall', 'f1' — so 'accuracy' is correct
fabric = default_fabric(monitor='accuracy', less_is_better=False)
```

**Tip**: The metric names come from `clf_metrics_builder` keys. Default names: `accuracy`, `precision`, `recall`, `f1`, plus optional `auroc`, `specificity`, etc.

## #10 Not calling load_best_ckpt before test

**Symptom**: Test scores are worse than expected — they reflect the last epoch's model, not the best.

**Cause**: After `train_loop`, the model is in its last-epoch state. The best checkpoint was saved to disk but not loaded back.

**Wrong**:
```python
procedure.train_loop(ctx, train_loader, val_loader=val_loader)
scores = procedure.test(ctx, test_loader)  # tests LAST epoch model
```

**Correct**:
```python
procedure.train_loop(ctx, train_loader, val_loader=val_loader)
ctx.load_best_ckpt('replace')  # load best model
scores = procedure.test(ctx, test_loader)  # tests BEST model
```

## Best-practice quick reference

| Scenario | Recommended setup |
|----------|------------------|
| Binary EEG classification | `BCEWithLogitsLoss` + `before_metric='sigmoid'` + `clf_metrics_builder(binary=True)` |
| Multiclass classification | `CrossEntropyLoss` + `before_metric='softmax'` + `clf_metrics_builder(task='multiclass', num_classes=N)` |
| Regression | `MSELoss` + `before_metric=None` + `reg_metrics_builder()` |
| Custom multi-loss | Custom step + `scalar_metrics='auto'` |
| GAN | Tuple optimizer + `scalar_metrics=['g_loss', 'd_loss']` + `fn_batch2target` override |
| EDL | No fn_loss + `scalar_metrics='auto'` + EDL procedure |
| K-fold | New fabric + context per fold, `version=f'fold_{k}'` |
