This is current my training code: ```python """ 시계...

This is current my training code:

python
"""
시계열(5 min resample) → 수면품질(Q1–S3) 멀티태스크 분류 (PyTorch Lightning).

주요 개선 (v4 - Encoder Fine-tuning)
------------------------------------
1. PyTorch Lightning 구조 유지.
2. **인코더 마지막 2개 블록 Fine-tuning 적용 (UNFREEZE_LAYERS = 2).**
3. 낮은 인코더 학습률(ENCODER_LR=1e-6) 유지.
4. 나머지 설정(LR, 스케줄러, 헤드 등)은 이전 버전(v3)과 동일하게 유지하여 효과 비교.
"""

# ──────────────────────────────────────────────────────────────────────────────
# 1. 라이브러리
import os
import random
import warnings
import logging
from pathlib import Path
from typing import List, Dict, Tuple, Optional, Any

import numpy as np
import pandas as pd
from tqdm.auto import tqdm

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score
import torch.amp

import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor
from pytorch_lightning.loggers import TensorBoardLogger # or WandbLogger, etc.

# torch.set_float32_matmul_precision('medium') # uncomment for potential speedup on Tensor Cores

from momentfm import MOMENTPipeline
from momentfm.utils.masking import Masking
from momentfm.common import TASKS

# ──────────────────────────────────────────────────────────────────────────────
# 2. 하이퍼파라미터 & 설정 (LightningModule/DataModule로 전달될 값들)
N_CHANNELS               = 22
LABEL_NAMES              = ['Q1', 'Q2', 'Q3', 'S1', 'S2', 'S3']
NUM_CLASSES_PER_LABEL    = [2, 2, 2, 3, 2, 2]
REDUCTION                = 'mean'
EPOCHS                   = 60
BATCH_SIZE               = 32
HEAD_LR                  = 5e-5            # 헤드 학습률 (유지)
ENCODER_LR               = 1e-6            # 인코더 학습률 (유지, UNFREEZE_LAYERS > 0 일 때 적용됨)
MAX_LR_FACTOR            = 1.0             # OneCycleLR 설정 (유지)
PATIENCE                 = 5               # EarlyStopping Patience (유지)
GRAD_CLIP                = 1.0
SEED                     = 42
RESAMPLE_RULE            = '5min'
EXPECTED_SEQ_LEN         = 288
TRAIN_DATA_PATH          = 'train_timeseries.parquet'
LABEL_DATA_PATH          = 'ch2025_metrics_train.csv'
VAL_SPLIT_RATIO          = 0.2
SAVE_DIR                 = Path('./lightning_logs_v4') # 로그 디렉토리 변경
CHECKPOINT_DIR           = SAVE_DIR / 'checkpoints'
CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
# --- *** 변경된 부분 *** ---
UNFREEZE_LAYERS          = 2               # 인코더 마지막 2개 블록 fine-tuning
# --------------------------
HEAD_DROPOUT             = 0.2
NUM_WORKERS              = 0
WEIGHT_DECAY             = 1e-2
PRECISION                = 'bf16-mixed'

warnings.filterwarnings('ignore', category=UserWarning)
warnings.filterwarnings("ignore", ".*Consider increasing the value of the `num_workers` argument*")
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# ──────────────────────────────────────────────────────────────────────────────
# 3. Dataset (기존과 동일)
class SensorDataset(Dataset):
    def __init__(self, X, y):
        self.X = X.astype(np.float32)
        self.y = y.astype(np.int64)
    def __len__(self): return len(self.X)
    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

# ──────────────────────────────────────────────────────────────────────────────
# 4. 데이터 로딩 및 전처리 함수 (기존과 동일)
def load_and_preprocess(time_series_path: str,
                        label_path: str,
                        resample_rule: str,
                        expected_len: int,
                        n_channels: int,
                        label_names: List[str]):
    logging.info('Loading parquet/csv for DataModule...')
    ts_df    = pd.read_parquet(time_series_path)
    label_df = pd.read_csv(label_path)

    ts_df['timestamp']   = pd.to_datetime(ts_df['timestamp'])
    ts_df['lifelog_date'] = ts_df['timestamp'].dt.date
    label_df['lifelog_date'] = pd.to_datetime(label_df['lifelog_date']).dt.date

    sensor_cols = [c for c in ts_df.columns if c not in
                   ('timestamp', 'subject_id', 'lifelog_date')]

    if len(sensor_cols) < n_channels:
         logging.warning(f'Data channels ({len(sensor_cols)}) < N_CHANNELS ({n_channels}).')
    elif len(sensor_cols) > n_channels:
        logging.warning(f'Data channels ({len(sensor_cols)}) > N_CHANNELS ({n_channels}). Using first {n_channels}.')
        sensor_cols = sensor_cols[:n_channels]

    merged = pd.merge(
        ts_df, label_df[['subject_id', 'lifelog_date'] + label_names],
        on=['subject_id', 'lifelog_date'], how='inner'
    )

    windows, labels = [], []
    grouped = merged.groupby(['subject_id', 'lifelog_date'])
    for (_, _), g in tqdm(grouped, desc='Resampling & Padding (DataModule)', total=len(grouped), leave=False):
        y_row = g.iloc[0][label_names].values.astype(np.int64)
        g = g.set_index('timestamp')[sensor_cols]
        g = g.ffill().fillna(0)
        g = g.resample(resample_rule).mean().fillna(0)

        current_len = len(g)
        if current_len < expected_len:
            pad_len = expected_len - current_len
            pad_start_time = g.index[-1] + pd.Timedelta(resample_rule) if not g.empty else pd.Timestamp.now()
            pad_index = pd.date_range(start=pad_start_time, periods=pad_len, freq=resample_rule)
            pad = pd.DataFrame(0, index=pad_index, columns=sensor_cols)
            g = pd.concat([g, pad])
        elif current_len > expected_len:
            g = g.iloc[-expected_len:]

        windows.append(g.values.T)
        labels.append(y_row)

    final_X = np.asarray(windows)
    final_y = np.asarray(labels)

    if final_y.dtype != np.int64:
        logging.warning(f"Final label array dtype was {final_y.dtype}, converting to int64.")
        final_y = final_y.astype(np.int64)

    return final_X, final_y

# ──────────────────────────────────────────────────────────────────────────────
# 5. LightningDataModule (기존과 동일)
class SleepDataModule(pl.LightningDataModule):
    def __init__(self,
                 data_path: str = TRAIN_DATA_PATH,
                 label_path: str = LABEL_DATA_PATH,
                 resample_rule: str = RESAMPLE_RULE,
                 expected_len: int = EXPECTED_SEQ_LEN,
                 n_channels: int = N_CHANNELS,
                 label_names: List[str] = LABEL_NAMES,
                 num_classes_per_label: List[int] = NUM_CLASSES_PER_LABEL,
                 batch_size: int = BATCH_SIZE,
                 val_split_ratio: float = VAL_SPLIT_RATIO,
                 seed: int = SEED,
                 num_workers: int = NUM_WORKERS):
        super().__init__()
        self.save_hyperparameters()
        self.X_train, self.y_train = None, None
        self.X_val, self.y_val = None, None
        self.task_class_weights: Optional[List[torch.Tensor]] = None

    def setup(self, stage: Optional[str] = None):
        if self.X_train is None and self.X_val is None:
            X, y = load_and_preprocess(
                self.hparams.data_path, self.hparams.label_path, self.hparams.resample_rule,
                self.hparams.expected_len, self.hparams.n_channels, self.hparams.label_names
            )
            logging.info(f"Raw data loaded: X shape {X.shape}, y shape {y.shape}")

            idx = np.arange(len(X))
            try:
                stratify_labels = y[:, 3]
                tr_idx, va_idx = train_test_split(idx, test_size=self.hparams.val_split_ratio,
                                                  random_state=self.hparams.seed, stratify=stratify_labels)
            except ValueError as e:
                 logging.warning(f"Stratification failed ({e}). Falling back to non-stratified split.")
                 tr_idx, va_idx = train_test_split(idx, test_size=self.hparams.val_split_ratio, random_state=self.hparams.seed)

            self.X_train, self.y_train = X[tr_idx], y[tr_idx]
            self.X_val, self.y_val = X[va_idx], y[va_idx]
            logging.info(f"Data split: Train {len(self.X_train)}, Validation {len(self.X_val)}")
            self.task_class_weights = self._calculate_class_weights(self.y_train)

    def _calculate_class_weights(self, train_labels_np: np.ndarray) -> List[torch.Tensor]:
        task_weights = []
        logging.info("Calculating class weights for each task...")
        for i, (name, num_classes) in enumerate(zip(self.hparams.label_names, self.hparams.num_classes_per_label)):
            labels_task_i = train_labels_np[:, i]
            counts = np.bincount(labels_task_i, minlength=num_classes)
            if np.any(counts == 0):
                logging.warning(f"Task '{name}': Found {np.sum(counts==0)} classes with zero samples in train set. Setting weight to 0.")
            weights = 1. / np.maximum(counts, 1)
            weights[counts == 0] = 0
            task_weights.append(torch.tensor(weights, dtype=torch.float32))
            logging.info(f"  Task '{name}' weights: {[f'{w:.3f}' for w in weights]}")
        return task_weights

    def train_dataloader(self):
        train_ds = SensorDataset(self.X_train, self.y_train)
        return DataLoader(train_ds, batch_size=self.hparams.batch_size, shuffle=True,
                          num_workers=self.hparams.num_workers, pin_memory=True, persistent_workers=bool(self.hparams.num_workers > 0))

    def val_dataloader(self):
        val_ds = SensorDataset(self.X_val, self.y_val)
        return DataLoader(val_ds, batch_size=self.hparams.batch_size, shuffle=False,
                          num_workers=self.hparams.num_workers, pin_memory=True, persistent_workers=bool(self.hparams.num_workers > 0))

# ──────────────────────────────────────────────────────────────────────────────
# 6. 경량 Multi‑Task Head (기존과 동일)
class SharedBottleneckHead(nn.Module):
    def __init__(self, in_dim: int, num_classes_per_label: List[int], dropout_rate: float):
        super().__init__()
        self.base = nn.Sequential(
            nn.Linear(in_dim, 512), nn.GELU(), nn.Dropout(dropout_rate),
            nn.Linear(512, 256), nn.GELU()
        )
        self.task_heads = nn.ModuleList([
            nn.Linear(256, c) for c in num_classes_per_label
        ])
        self.label_names = LABEL_NAMES

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        shared = self.base(x)
        return {name: head(shared) for name, head in zip(self.label_names, self.task_heads)}

# ──────────────────────────────────────────────────────────────────────────────
# 7. LightningModule (unfreeze_layers 적용 부분 외 기존과 동일)
class SleepQualityMoment(pl.LightningModule):
    def __init__(self,
                 n_channels: int = N_CHANNELS,
                 label_names: List[str] = LABEL_NAMES,
                 num_classes_per_label: List[int] = NUM_CLASSES_PER_LABEL,
                 task_class_weights: List[torch.Tensor] = None,
                 reduction: str = REDUCTION,
                 unfreeze_layers: int = UNFREEZE_LAYERS, # 이 값이 이제 2로 전달됨
                 head_dropout: float = HEAD_DROPOUT,
                 lr_head: float = HEAD_LR,
                 lr_encoder: float = ENCODER_LR,
                 max_lr_factor: float = MAX_LR_FACTOR,
                 weight_decay: float = WEIGHT_DECAY,
                 total_steps: Optional[int] = None):
        super().__init__()
        self.save_hyperparameters(ignore=['task_class_weights'])
        self.task_class_weights = task_class_weights
        self._build_model()
        self.criterion = nn.CrossEntropyLoss

    def _build_model(self):
        # MOMENT 로드 및 기본 설정 (기존과 동일)
        self.pipe_base = MOMENTPipeline.from_pretrained(
            "AutonLab/MOMENT-1-large",
            model_kwargs=dict(
                task_name=TASKS.EMBED, n_channels=self.hparams.n_channels,
                freeze_encoder=True, freeze_embedder=True,
                reduction=self.hparams.reduction,
            )
        )
        self.normalizer = self.pipe_base.normalizer
        self.tokenizer = self.pipe_base.tokenizer
        self.patch_embedding = self.pipe_base.patch_embedding
        self.encoder = self.pipe_base.encoder
        self.d_model = self.pipe_base.config.d_model
        self.patch_len = self.pipe_base.config.patch_len
        self.reduction = self.hparams.reduction

        for param in self.parameters():
            param.requires_grad = False

        # --- 인코더 Unfreeze 로직 (이제 unfreeze_layers=2 적용) ---
        if self.hparams.unfreeze_layers > 0:
            logging.info(f"Unfreezing last {self.hparams.unfreeze_layers} encoder blocks.")
            enc_blocks = self._get_encoder_blocks()
            num_total_blocks = len(enc_blocks)
            if num_total_blocks < self.hparams.unfreeze_layers:
                 logging.warning(f"Requested unfreeze ({self.hparams.unfreeze_layers}) > available ({num_total_blocks}). Unfreezing all.")
                 self.hparams.unfreeze_layers = num_total_blocks

            for i in range(num_total_blocks - self.hparams.unfreeze_layers, num_total_blocks):
                for param in enc_blocks[i].parameters():
                    param.requires_grad = True
        else:
            logging.info("Keeping encoder fully frozen.")
        # ----------------------------------------------------

        # 헤드 정의 및 활성화 (기존과 동일)
        in_dim = self.d_model if self.reduction == "mean" else self.d_model * self.hparams.n_channels
        self.head = SharedBottleneckHead(in_dim, self.hparams.num_classes_per_label, self.hparams.head_dropout)
        for param in self.head.parameters():
            param.requires_grad = True

        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
        total_params = sum(p.numel() for p in self.parameters())
        logging.info(f"Model built: Trainable params: {trainable_params:,} / Total params: {total_params:,}")


    def _get_encoder_blocks(self):
        for attr in ("block", "layer", "layers"):
            if hasattr(self.encoder, attr):
                return getattr(self.encoder, attr)
        raise AttributeError("Cannot find encoder blocks.")

    def get_representation(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        # 기존과 동일
        b, c, t = x.shape
        x = self.normalizer(x=x, mask=mask, mode='norm')
        x = torch.nan_to_num(x)
        x_tok = self.tokenizer(x=x)
        emb = self.patch_embedding(x_tok, mask=mask)
        n_patch, d_model = emb.shape[2], emb.shape[3]
        emb = emb.reshape(b * c, n_patch, d_model)
        patch_mask = Masking.convert_seq_to_patch_view(mask, self.patch_len)
        attn_mask = patch_mask.repeat_interleave(c, 0)
        enc_out = self.encoder(inputs_embeds=emb, attention_mask=attn_mask).last_hidden_state
        enc_out = enc_out.reshape(b, c, n_patch, d_model)
        if self.reduction == 'mean':
            enc_out = enc_out.mean(dim=1)
        elif self.reduction == 'concat':
            enc_out = enc_out.permute(0, 2, 3, 1).reshape(b, n_patch, d_model * c)
        else:
            enc_out = enc_out.mean(dim=1)
        patch_mask = patch_mask.unsqueeze(-1)
        rep = (enc_out * patch_mask).sum(dim=1) / patch_mask.sum(dim=1).clamp_min(1e-5)
        return rep

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> Dict[str, torch.Tensor]:
        # 기존과 동일
        rep = self.get_representation(x, mask)
        logits = self.head(rep)
        return logits

    def _calculate_loss(self, logits: Dict[str, torch.Tensor], targets: torch.Tensor) -> torch.Tensor:
        # 기존과 동일
        total_loss = 0.0
        for i, name in enumerate(self.hparams.label_names):
            weights = self.task_class_weights[i].to(logits[name].device)
            loss_fn_task = self.criterion(weight=weights)
            task_loss = loss_fn_task(logits[name], targets[:, i])
            total_loss += task_loss
        return total_loss / len(self.hparams.label_names)

    def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:
        # 기존과 동일
        x, y = batch
        mask = torch.ones((x.size(0), x.size(2)), dtype=torch.bool, device=x.device)
        logits = self(x, mask)
        loss = self._calculate_loss(logits, y)
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
        return loss

    def validation_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int):
        # 기존과 동일
        x, y = batch
        mask = torch.ones((x.size(0), x.size(2)), dtype=torch.bool, device=x.device)
        logits = self(x, mask)
        loss = self._calculate_loss(logits, y)
        avg_f1, task_f1 = self._compute_f1(logits, y)
        self.log('val_loss', loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)
        self.log('val_f1_macro', avg_f1, on_step=False, on_epoch=True, prog_bar=True, logger=True)
        for name, f1_val in task_f1.items():
            self.log(f'val_f1_{name}', f1_val, on_step=False, on_epoch=True, logger=True)

    def _compute_f1(self, pred_logits: Dict[str, torch.Tensor], true: torch.Tensor) -> Tuple[float, Dict[str, float]]:
        # 기존과 동일
        scores = {}
        true_np = true.cpu().numpy()
        for i, name in enumerate(self.hparams.label_names):
            pred_np = torch.argmax(pred_logits[name], 1).cpu().numpy()
            f1 = f1_score(true_np[:, i], pred_np, average='macro', zero_division=0)
            scores[name] = f1
        average_f1 = float(np.mean(list(scores.values())))
        return average_f1, scores

    def configure_optimizers(self) -> Dict[str, Any]:
        # 옵티마이저 설정 (기존과 동일, unfreeze_layers > 0 이므로 인코더 파라미터 그룹 추가됨)
        param_groups = [
            {"params": self.head.parameters(), "lr": self.hparams.lr_head, "name": "Head"},
        ]
        encoder_params = [p for p in self.encoder.parameters() if p.requires_grad]
        if encoder_params:
            param_groups.append({"params": encoder_params, "lr": self.hparams.lr_encoder, "name": "EncoderFinetune"})
            logging.info(f"Encoder params added to optimizer with LR: {self.hparams.lr_encoder}")
        else:
             logging.warning("No encoder parameters require grad, check unfreeze_layers setting.")


        optimizer = optim.AdamW(param_groups, weight_decay=self.hparams.weight_decay)

        if self.hparams.total_steps is None:
             logging.warning("total_steps not provided to LightningModule, scheduler might not be optimal.")
             total_steps = 1000000 # Fallback, should be calculated in main script
        else:
            total_steps = self.hparams.total_steps

        scheduler = optim.lr_scheduler.OneCycleLR(
            optimizer,
            max_lr=[pg.get('lr', self.hparams.lr_head) * self.hparams.max_lr_factor for pg in param_groups],
            total_steps=total_steps,
            pct_start=0.1,
        )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": 1
            },
        }

# ==============================================================================
# Main Execution Block (기존과 동일, 변경된 UNFREEZE_LAYERS 값이 전달됨)
# ==============================================================================
if __name__ == '__main__':
    pl.seed_everything(SEED, workers=True)

    # 1. Initialize DataModule
    dm = SleepDataModule(
        data_path=TRAIN_DATA_PATH, label_path=LABEL_DATA_PATH,
        resample_rule=RESAMPLE_RULE, expected_len=EXPECTED_SEQ_LEN,
        n_channels=N_CHANNELS, label_names=LABEL_NAMES,
        num_classes_per_label=NUM_CLASSES_PER_LABEL, batch_size=BATCH_SIZE,
        val_split_ratio=VAL_SPLIT_RATIO, seed=SEED, num_workers=NUM_WORKERS
    )
    dm.setup('fit')

    # 2. Calculate total_steps for Scheduler
    steps_per_epoch = len(dm.train_dataloader())
    # Check if steps_per_epoch is valid
    if steps_per_epoch == 0:
        raise ValueError("Train DataLoader is empty. Check dataset or batch size.")
    total_training_steps = steps_per_epoch * EPOCHS
    logging.info(f"Calculated total training steps: {total_training_steps}")

    # 3. Initialize LightningModule (이제 UNFREEZE_LAYERS=2 로 초기화됨)
    model = SleepQualityMoment(
        n_channels=N_CHANNELS, label_names=LABEL_NAMES,
        num_classes_per_label=NUM_CLASSES_PER_LABEL,
        task_class_weights=dm.task_class_weights,
        reduction=REDUCTION,
        unfreeze_layers=UNFREEZE_LAYERS, # *** 여기가 핵심 변경 ***
        head_dropout=HEAD_DROPOUT,
        lr_head=HEAD_LR, lr_encoder=ENCODER_LR,
        max_lr_factor=MAX_LR_FACTOR, weight_decay=WEIGHT_DECAY,
        total_steps=total_training_steps
    )

    # 4. Initialize Callbacks (기존과 동일)
    checkpoint_callback = ModelCheckpoint(
        dirpath=CHECKPOINT_DIR,
        filename='best-model-{epoch:02d}-{val_f1_macro:.4f}',
        monitor='val_f1_macro', mode='max', save_top_k=1, save_last=False, verbose=True
    )
    early_stopping_callback = EarlyStopping(
        monitor='val_f1_macro', mode='max', patience=PATIENCE, verbose=True
    )
    lr_monitor = LearningRateMonitor(logging_interval='step')

    # 5. Initialize Logger (기존과 동일)
    logger = TensorBoardLogger(SAVE_DIR, name="sleep_moment_finetune_v4") # Logger 이름 변경

    # 6. Initialize Trainer (기존과 동일)
    trainer = pl.Trainer(
        max_epochs=EPOCHS, accelerator="auto", devices="auto",
        precision=PRECISION, gradient_clip_val=GRAD_CLIP,
        logger=logger,
        callbacks=[checkpoint_callback, early_stopping_callback, lr_monitor],
        deterministic=False,
        log_every_n_steps=min(50, steps_per_epoch),
    )

    # 7. Start Training
    logging.info("Starting training with PyTorch Lightning (Encoder Fine-tuning Enabled)...")
    trainer.fit(model, datamodule=dm)

    logging.info("Training finished.")
    logging.info(f"Best model checkpoint saved at: {checkpoint_callback.best_model_path}")

    # Optional: Validate the best model after training
    logging.info("Validating the best model...")
    val_results = trainer.validate(model, datamodule=dm, ckpt_path='best')
    logging.info(f"Validation results of the best model: {val_results}")

This is my current Log:

text
Seed set to 42
2025-05-02 08:24:52,972 - INFO - Loading parquet/csv for DataModule...
2025-05-02 08:26:06,100 - INFO - Raw data loaded: X shape (450, 22, 288), y shape (450, 6)
2025-05-02 08:26:06,107 - INFO - Data split: Train 360, Validation 90
2025-05-02 08:26:06,107 - INFO - Calculating class weights for each task...
2025-05-02 08:26:06,108 - INFO -   Task 'Q1' weights: ['0.006', '0.006']
2025-05-02 08:26:06,109 - INFO -   Task 'Q2' weights: ['0.006', '0.005']
2025-05-02 08:26:06,109 - INFO -   Task 'Q3' weights: ['0.007', '0.005']
2025-05-02 08:26:06,109 - INFO -   Task 'S1' weights: ['0.009', '0.006', '0.015']
2025-05-02 08:26:06,109 - INFO -   Task 'S2' weights: ['0.008', '0.004']
2025-05-02 08:26:06,110 - INFO -   Task 'S3' weights: ['0.008', '0.004']
2025-05-02 08:26:06,112 - INFO - Calculated total training steps: 720
2025-05-02 08:26:06,227 - INFO - Setting d_model to 1024
2025-05-02 08:26:09,199 - INFO - Initializing pre-trained transformer from google/flan-t5-large.
2025-05-02 08:26:09,201 - INFO - Enabling gradient checkpointing.
2025-05-02 08:26:09,938 - INFO - Unfreezing last 2 encoder blocks.
2025-05-02 08:26:09,943 - INFO - Model built: Trainable params: 26,353,677 / Total params: 341,907,989
Using bfloat16 Automatic Mixed Precision (AMP)
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
2025-05-02 08:26:09,953 - INFO - Starting training with PyTorch Lightning (Encoder Fine-tuning Enabled)...
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
2025-05-02 08:26:10,226 - INFO - Encoder params added to optimizer with LR: 1e-06

  | Name            | Type                 | Params | Mode 
-----------------------------------------------------------------
0 | pipe_base       | MOMENTPipeline       | 341 M  | train
1 | normalizer      | RevIN                | 0      | train
2 | tokenizer       | Patching             | 0      | train
3 | patch_embedding | PatchEmbedding       | 9.2 K  | train
4 | encoder         | T5Stack              | 341 M  | train
5 | head            | SharedBottleneckHead | 659 K  | train
-----------------------------------------------------------------
26.4 M    Trainable params
315 M     Non-trainable params
341 M     Total params
1,367.632 Total estimated model params size (MB)
486       Modules in train mode
0         Modules in eval mode
Epoch 8: 100%
 12/12 [00:07<00:00,  1.67it/s, v_num=0, train_loss_step=0.761, val_loss=0.761, val_f1_macro=0.300, train_loss_epoch=0.761]
Metric val_f1_macro improved. New best score: 0.271
Epoch 0, global step 12: 'val_f1_macro' reached 0.27119 (best 0.27119), saving model to '/workspace/lightning_logs_v4/checkpoints/best-model-epoch=00-val_f1_macro=0.2712.ckpt' as top 1
Epoch 1, global step 24: 'val_f1_macro' was not in top 1
Epoch 2, global step 36: 'val_f1_macro' was not in top 1
Metric val_f1_macro improved by 0.048 >= min_delta = 0.0. New best score: 0.319
Epoch 3, global step 48: 'val_f1_macro' reached 0.31931 (best 0.31931), saving model to '/workspace/lightning_logs_v4/checkpoints/best-model-epoch=03-val_f1_macro=0.3193.ckpt' as top 1
Epoch 4, global step 60: 'val_f1_macro' was not in top 1
Epoch 5, global step 72: 'val_f1_macro' was not in top 1
Epoch 6, global step 84: 'val_f1_macro' was not in top 1
Epoch 7, global step 96: 'val_f1_macro' was not in top 1
Monitored metric val_f1_macro did not improve in the last 5 records. Best score: 0.319. Signaling Trainer to stop.
Epoch 8, global step 108: 'val_f1_macro' was not in top 1
2025-05-02 08:27:23,805 - INFO - Training finished.
2025-05-02 08:27:23,805 - INFO - Best model checkpoint saved at: /workspace/lightning_logs_v4/checkpoints/best-model-epoch=03-val_f1_macro=0.3193.ckpt
2025-05-02 08:27:23,806 - INFO - Validating the best model...
Restoring states from the checkpoint path at /workspace/lightning_logs_v4/checkpoints/best-model-epoch=03-val_f1_macro=0.3193.ckpt
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Loaded model weights from the checkpoint at /workspace/lightning_logs_v4/checkpoints/best-model-epoch=03-val_f1_macro=0.3193.ckpt
Validation DataLoader 0: 100%
 3/3 [00:01<00:00,  2.19it/s]
────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
     Validate metric           DataLoader 0
────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
        val_f1_Q1           0.3385222256183624
        val_f1_Q2           0.2811170220375061
        val_f1_Q3           0.3626408874988556
        val_f1_S1           0.16120068728923798
        val_f1_S2           0.5245703458786011
        val_f1_S3           0.2478388100862503
      val_f1_macro          0.31931498646736145
        val_loss            0.7615204453468323
────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
2025-05-02 08:27:27,860 - INFO - Validation results of the best model: [{'val_loss': 0.7615204453468323, 'val_f1_macro': 0.31931498646736145, 'val_f1_Q1': 0.3385222256183624, 'val_f1_Q2': 0.2811170220375061, 'val_f1_Q3': 0.3626408874988556, 'val_f1_S1': 0.16120068728923798, 'val_f1_S2': 0.5245703458786011, 'val_f1_S3': 0.2478388100862503}]
Selection deleted

My current f1 score is too low. Please rewrite the entire code by applying advanced Sota techniques so that the verification f1 score can reach 0.8 or higher.