exp_timeseries.py

import torch
import pytorch_lightning as pl

from collections import defaultdict
from models.Basic import MLP
from models.Lstm import LSTM
from models.Informer import Informer, InformerStack
from models.Stockformer import Stockformer
from utils.stock_metrics import get_stock_algo, pct_direction_torch
from torchmetrics import MeanSquaredError, MeanAbsoluteError
from torch_optimizer import Ranger


class ExpTimeseries(pl.LightningModule):
    def __init__(self, config):
        super().__init__()
        self.config = config

        # pl makes self.learning_rate special
        self.learning_rate = config.learning_rate

        # Torch metrics has a state that resets but val and train can be called in unison so we split
        # If pre_loss isn't supplied (ie: pre_loss is None) it will default to config.loss
        self.train_criterion = self._select_criterion(
            loss_override=self.config.pre_loss
        )
        self.other_criterion = self._select_criterion(
            loss_override=self.config.pre_loss
        )
        self.metric = self._select_criterion(metric=True)
        self.loss_switched = False

        self._build_model()
        # self.save_hyperparameters()
        self.loss_reg = None
        self.scale = None
        self.val_log_growth_sum = None
        self.val_log_growth_count = None
        self.test_log_growth_sum = None
        self.test_log_growth_count = None

    def _build_model(self):
        model_dict = {
            "informer": Informer,
            "informerstack": InformerStack,
            "mlp": MLP,
            "stockformer": Stockformer,
            "lstm": LSTM,
        }
        assert (
            self.config.model in model_dict
        ), f"Invalid config.model: {self.config.model}, options: {list(model_dict.keys())}"
        self.model = model_dict[self.config.model](self.config).float()

        # Load model
        if self.config.load_model_path is not None:
            self.load_from_checkpoint(self.config.load_model_path)

    def _select_criterion(self, loss_override=None, metric=False):
        loss = self.config.loss
        if loss_override is not None:
            loss = loss_override

        def combine_loss(loss, weights=None):
            if weights is None:
                weights = [1.0] * len(loss)
            def combined(pred, target, inv_pred, input_scale=None):
                return loss[0](pred, target, input_scale=input_scale)
                #  return sum(w*l(inv_pred, target) if "Mean" in l.__class__.__name__ else w*l(pred, target) for w,l in zip(weights, loss))
            return combined

        def loss_lib(loss: str):
            if "stock" in loss:
                # Using Stock Loss
                _, stock_loss_mode = loss.split("_")
                target_type = self.config.target.split("_")[1]
                assert (
                    target_type == "pctchange" or target_type == "logpctchange"
                ), "Can't use stock loss unless target is pctchange or logpctchange"
                assert (
                    self.config.scale and
                    self.config.inverse_pred
                    #  and not self.config.inverse_output
                ), "Can't use stock loss without scale, inverse pred, and not inverse output"

                criterion = get_stock_algo(target_type, stock_loss_mode)
                print("criterion:", criterion)
                if metric:
                    def mt(x, y, input_scale):
                        return criterion.metric(x, y, input_scale=input_scale)
                    return mt
                else:
                    return lambda x, y, input_scale: [-1 * criterion.loss(x, y).mean(), criterion.sharpe(x, y).mean()]
                # return lambda x, y: -LogPctProfitTanhV1.loss(x, y).mean()
                # return get_stock_loss(target_type, stock_loss_mode, threshold=0.0)
            elif loss == "mae":
                assert (
                    self.config.scale
                    and self.config.inverse_pred
                    #  and self.config.inverse_output
                ), "Can't use mae loss without scale, inverse pred, and inverse output"
                return MeanAbsoluteError().cuda()
            elif loss == "mse":
                assert (
                    self.config.scale
                    and self.config.inverse_pred
                    #  and self.config.inverse_output
                ), "Can't use mse loss without scale, inverse pred, and inverse output"
                return MeanSquaredError().cuda()
        loss_list = [ loss_lib(loss_type) for loss_type in loss.split('+') ]
        weights = [1.0] if '+' not in loss else [1.0, 0.1]
        return combine_loss(loss_list, weights)

        raise Exception(f"Invalid loss: {loss}")

    def forward(self, x):
        # in lightning, forward defines the prediction/inference actions
        return self.model(x)

    def training_step(self, batch, batch_idx):
        # training_step defines the train loop. It is independent of forward
        batch_x, batch_y, batch_x_mark, batch_y_mark, _ = batch

        sigma_x = 0.001
        batch_x = batch_x + (torch.randn_like(batch_x)*2-1) * sigma_x
        #  print(sigma_x.mean(), batch_x.mean(), batch_x.shape)
        #  sigma_y = 0.01 * batch_y.std(dim=(1, 2), keepdim=True)
        #  batch_y = batch_y + (torch.randn_like(batch_y)*2-1) * sigma_y

        pred, true, inv_pred = self._process_one_batch(
            self.trainer.datamodule.data_train,
            batch_x,
            batch_y,
            batch_x_mark,
            batch_y_mark,
            ds_index=None,
        )
        #  print(self.loss_reg)
        loss, sharpe = self.train_criterion(pred, true, inv_pred)
        self.log("train_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
        self.log("train_sharpe", sharpe, prog_bar=True, on_step=False, on_epoch=True)
        self.log("wavelet_loss", self.loss_reg, prog_bar=True, on_step=False, on_epoch=True)

        #  self.log(
            #  "tr_pct_dir",
            #  pct_direction_torch(pred, true),
            #  prog_bar=True,
            #  on_step=False,
            #  on_epoch=True,
        #  )
        #  self.log(
            #  "tr_mag",
            #  torch.linalg.norm(pred),  # torch.mean(torch.abs(pred))
            #  prog_bar=False,
            #  on_step=False,
            #  on_epoch=True,
        #  )

        if (
            self.config.pre_epochs is not None
            and self.config.pre_loss is not None
            and self.current_epoch == self.config.pre_epochs
            and not self.loss_switched
        ):
            # Revert to default loss
            self.train_criterion = self._select_criterion(
                loss_override=self.config.loss
            )
            self.other_criterion = self._select_criterion(
                loss_override=self.config.loss
            )
            self.loss_switched = True

        return loss + torch.exp(-2.5*sharpe) + 1e0*self.loss_reg

    def validation_step(self, batch, batch_idx, dataloader_idx=0):
        # validation_step defines the validation loop. It is independent of forward
        batch_x, batch_y, batch_x_mark, batch_y_mark, _ = batch

        pred, true, inv_pred = self._process_one_batch(
            self.trainer.datamodule.data_val,
            batch_x,
            batch_y,
            batch_x_mark,
            batch_y_mark,
            ds_index=None,
        )

        if dataloader_idx == 0:
            # Actual val dataset
            #  assert self.trainer.val_dataloaders[0].dataset.flag == "val"
            loss, sharpe = self.other_criterion(pred, true, inv_pred)
            self.log(
                "val_loss",
                loss,
                prog_bar=True,
                on_step=False,
                on_epoch=True,
                sync_dist=False,
                add_dataloader_idx=False,
            )
            self.log(
                "val_sharpe",
                sharpe,
                prog_bar=True,
                on_step=False,
                on_epoch=True,
                sync_dist=False,
                add_dataloader_idx=False,
            )
            raw, self.scale = self.metric(pred, true, inv_pred)
            self.val_log_growth_sum[0] += raw.detach().sum()
            self.val_log_growth_count[0] += raw.numel()
            #  self.log(
                #  "val_pct_dir",
                #  pct_direction_torch(pred, true),
                #  prog_bar=False,
                #  on_step=False,
                #  on_epoch=True,
                #  add_dataloader_idx=False,
            #  )
            return
        elif dataloader_idx == 1:
            # TODO: If we are using torch metrics we should create an additional loss function
            # Test dataset
            assert self.trainer.val_dataloaders[1].dataset.flag == "test"
            loss, sharpe = self.other_criterion(pred, true, inv_pred)
            self.log(
                "test_loss",
                loss,
                prog_bar=True,
                on_step=False,
                on_epoch=True,
                sync_dist=False,
                add_dataloader_idx=False,
            )
            self.log(
                "test_sharpe",
                sharpe,
                prog_bar=True,
                on_step=False,
                on_epoch=True,
                sync_dist=False,
                add_dataloader_idx=False,
            )
            raw, _ = self.metric(pred, true, inv_pred, self.scale)
            self.val_log_growth_sum[1] += raw.detach().sum()
            self.val_log_growth_count[1] += raw.numel()
            #  self.log(
                #  "test_pct_dir",
                #  pct_direction_torch(pred, true),
                #  prog_bar=False,
                #  on_step=False,
                #  on_epoch=True,
                #  add_dataloader_idx=False,
            #  )
            return

    def on_validation_epoch_start(self):
        self.val_log_growth_sum = defaultdict(lambda: 0.0)
        self.val_log_growth_count = defaultdict(int)

    def on_validation_epoch_end(self):
        for dl_idx, sum_log in self.val_log_growth_sum.items():
            #  count = self.val_log_growth_count[dl_idx]
            factor = torch.exp(sum_log)
            roi = factor - 1
            if dl_idx == 0:
                name = "val_roi"
            elif dl_idx == 1:
                name = "test_roi"
            else:
                raise Exception
            self.log(
                name,
                roi,
                prog_bar=True,
                on_step=False,
                on_epoch=True,
                sync_dist=False,
                add_dataloader_idx=False,
            )
        #  # 或者用平均 log-growth 当 metric（和 T 无关，更稳）
        #  mean_log_growth = self.val_log_growth_sum / self.val_log_growth_count
        #  self.log("val_mean_log_growth", mean_log_growth, prog_bar=False)

    def test_step(self, batch, batch_idx, dataloader_idx=0):
        # test_step defines the test loop. It is independent of forward
        batch_x, batch_y, batch_x_mark, batch_y_mark, _ = batch

        data_sets = [
            self.trainer.datamodule.data_train,
            self.trainer.datamodule.data_val,
            self.trainer.datamodule.data_test,
        ]

        pred, true, inv_pred = self._process_one_batch(
            data_sets[dataloader_idx],
            batch_x,
            batch_y,
            batch_x_mark,
            batch_y_mark,
            ds_index=None,
        )
        #  loss = self.other_criterion(pred, true, inv_pred)
        #  # if dataloader_idx == 0:
        #  self.log(
            #  "test_loss",
            #  loss,
            #  sync_dist=False,
        #  )

        if dataloader_idx == 0:
            raw, _ = self.metric(pred, true, inv_pred)
        if dataloader_idx == 1:
            raw, self.scale = self.metric(pred, true, inv_pred)
        if dataloader_idx == 2:
            raw, _ = self.metric(pred, true, inv_pred, self.scale)
        self.test_log_growth_sum[dataloader_idx] += raw.detach().sum()
        self.test_log_growth_count[dataloader_idx] += raw.numel()

    def on_test_epoch_start(self):
        self.test_log_growth_sum = defaultdict(lambda: 0.0)
        self.test_log_growth_count = defaultdict(int)

    def on_test_epoch_end(self):
        for dl_idx, sum_log in self.test_log_growth_sum.items():
            factor = torch.exp(sum_log)
            roi = factor - 1
            self.log(
                "test_roi",
                roi,
                sync_dist=False,
            )

    def predict_step(self, batch, batch_idx, dataloader_idx=0):
        batch_x, batch_y, batch_x_mark, batch_y_mark, _ = batch

        data_sets = [
            self.trainer.datamodule.data_train,
            self.trainer.datamodule.data_val,
            self.trainer.datamodule.data_test,
        ]

        pred, true, inv_pred = self._process_one_batch(
            data_sets[dataloader_idx],
            batch_x,
            batch_y,
            batch_x_mark,
            batch_y_mark,
            ds_index=None,
        )

        # dataset = self.trainer.predict_dataloaders[dataloader_idx].dataset
        # batch_x_raw_date, batch_y_raw_date = dataset.index_to_dates(batch_idx)

        if "mse" in self.config.loss or "mae" in self.config.loss:
            pred = inv_pred
        return {
            "pred": pred.detach().to(torch.float32),
            "true": true.detach().to(torch.float32),
        }

    # def on_predict_epoch_end(self, results):
    #     pass

    # def on_predict_end(self):
    #     pass

    def _process_one_batch(
        self,
        dataset_object,
        batch_x,
        batch_y,
        batch_x_mark,
        batch_y_mark,
        ds_index=None,
    ):
        # Decoder input if self.config.dec_in
        dec_inp = None
        # if self.config.dec_in and (
        #     self.config.padding == 0 or self.config.padding == 1
        # ):
        #     # FF: dec_inp = torch.zeros_like(batch_y[:, -self.config.pred_len:, :]).float()
        #     dec_inp = torch.full(
        #         [batch_y.shape[0], self.config.pred_len, batch_y.shape[-1]],
        #         self.config.padding,
        #     ).float()
        #     dec_inp = (
        #         torch.cat([batch_y[:, : self.config.label_len, :], dec_inp], dim=1)
        #         .float()
        #         .to(self.device)
        #     )

        # Encoder - Decoder
        if self.config.output_attention:
            outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
        else:
            outputs, loss_reg = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
            self.loss_reg = loss_reg
        #  if self.config.inverse_output:
        f_dim = -1 if self.config.features == "MS" else 0

        # if ds_index is None:
        batch_y = batch_y[:, -self.config.pred_len :, f_dim:]
        #  print(batch_y.std())
        #  batch_y = dataset_object.inverse_transform(batch_y)
        #  print(batch_y.std())
        #  while 1:pass
        inv_outputs = dataset_object.inverse_transform(outputs)
        return outputs, batch_y, inv_outputs
        # else:
        #     batch_x_raw_dates, batch_y_raw_dates = dataset_object.index_to_dates(
        #         ds_index
        #     )
        #     assert batch_y_raw_dates.shape == batch_y.shape[0:2]
        #     batch_y = batch_y[:, -self.config.pred_len :, f_dim:].to(self.device)
        #     batch_y_raw_dates = batch_y_raw_dates[:, -self.config.pred_len :]
        #     return outputs, batch_y, batch_y_raw_dates

    def configure_optimizers(self):
        if self.config.optim == "AdamW":
            optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate)
        elif self.config.optim == "Ranger":
            optimizer = Ranger(self.parameters(), lr=self.learning_rate)
        else:
            optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        # optimizer = torch.optim.SGD(self.parameters(), lr=self.learning_rate)

        # Learning rate scheduler
        if self.config.lradj == "type1":
            lmbda = lambda epoch: 0.5
            scheduler = torch.optim.lr_scheduler.MultiplicativeLR(
                optimizer, lr_lambda=lmbda, verbose=True
            )
        elif self.config.lradj == "type2":
            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                optimizer,
                factor=0.5,
                patience=10,
                threshold=0,
                cooldown=0,
                verbose=True,
                min_lr=1e-8,
            )
            scheduler = {
                "scheduler": scheduler,
                "interval": "epoch",  # called after each training epoch
                "monitor": "val_loss",
            }
        elif self.config.lradj == "type3":
            scheduler = torch.optim.lr_scheduler.OneCycleLR(
                optimizer,
                max_lr=self.config.learning_rate,
                steps_per_epoch=len(self.trainer.datamodule.data_train)
                // self.config.batch_size,  # Would be nicer to use self.trainer.train_dataloader.dataset but there is a pl bug
                epochs=self.config.max_epochs,
            )
            scheduler = {
                "scheduler": scheduler,
                "interval": "step",  # called after each training step
            }
        else:
            return optimizer

        return [optimizer], [scheduler]