# Regression Task

## Overview
[regression.py](https://github.com/ofsoundof/TimeFM/blob/add_documentation/tasks/regression_task.py) implements a fine-tuning task workflow for specifically for regression problems. The `RegressionTask` leverages [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/) for training/validation loops, along with data augmentations (white noise or spectrogram masking), regression-specific metrics (Pearson correlation, R2 score, RMSE), and optional layerwise learning rate decay for fine-tuning a backbone.

---

## Key Components

### 1. **Initialization**

- **Encoder (`self.model`)**: A backbone that processes the audio data (waveforms or spectrograms).
- **Regression Head (`self.model_head`)**: Produces continuous outputs, returning vectors whose size matches the regression target dimension.
- **Freeze Backbone**: If `freeze_backbone=True`, encoder parameters are frozen.
- **Augmentations**:
  - **WhiteNoiseAugment**: Adds Gaussian noise to the input.  
  - **SpecAugment**: Optionally applies frequency/time masking to spectrograms.
- **Metrics**:
  - **PearsonCorrCoef**: Measures linear correlation between predictions and targets.  
  - **R2Score**: Coefficient of determination (default: “uniform_average” for multioutput).
  - **MeanSquaredError** (with `squared=False` to get RMSE).

### 2. **Training Step**

1. **Data Flow**:  
   - `batch['input']` -> `X`  
   - `batch['label']` -> `y`  
   - (Optional) `batch['nr_padded_channels']` if padding is used.
2. **Forward Pass**:
   - `encoder_output = self.model(X)`
   - `y_preds = self.model_head(encoder_output)` -> squeezed + passed through `nn.Tanh()` to force output in `[-1, 1]`.
3. **Loss & Metrics**:
   - `loss = self.criterion(y_preds_flat, batch)` uses an MSE-based criterion.
   - `self.train_pearson`, `self.train_r2`, and `self.train_rmse` are updated with `y_preds_flat` and flattened ground truth.
4. **Logging**:
   - Training loss is logged as `train_loss` at both step and epoch levels.

### 3. **Validation & Test Steps**

- **Validation** (`validation_step`) and **Test** (`test_step`) have the same forward pass:
  1. Extract inputs and labels.
  2. Forward pass through encoder and regression head, apply `nn.Tanh()`.
  3. Compute loss, update metrics (`val_pearson`, `val_r2`, `val_rmse` or test equivalents).
  4. Log results (`val_loss` or `test_loss`) on step and epoch.

- **End-of-Epoch Logging**:  
  - `on_validation_epoch_end` logs `val_r2`.  
  - `on_test_epoch_end` logs `test_r2` and `test_rmse`.

### 4. **On-the-Fly Augmentations**

In `on_after_batch_transfer`, the script optionally:
1. Adds white noise if training (`self.white_noise_augment`).
2. Converts input to spectrogram (via `self.transform`) if `using_spectrogram=True`.
3. Applies SpecAugment frequency/time masking (`self.spec_augment`).

### 5. **Layerwise Learning Rate Decay**

In `configure_optimizers`:
- `num_blocks = self.hparams.model.depth` indicates the number of transformer layers.
- Each layer can receive a decayed learning rate (`base_lr * decay_factor^(num_blocks - block_nr)`).
- This can help stabilize fine-tuning deeper layers while leaving initial layers less frequently updated.

### 6. **Optimizer & Scheduler**

- The task supports multiple optimizers: `SGD`, `Adam`, `AdamW`, `LAMB`.
- Instantiates a scheduler from Hydra config:
  ```python
  scheduler = hydra.utils.instantiate(
      self.hparams.scheduler,
      optimizer=optimizer,
      total_training_opt_steps=self.trainer.estimated_stepping_batches
  )
