242 lines
7.1 KiB
Python
242 lines
7.1 KiB
Python
import argparse
|
|
import logging
|
|
import wandb
|
|
import pprint
|
|
import os
|
|
import pandas as pd
|
|
import lightning.pytorch as pl
|
|
|
|
from lightning.pytorch.utilities.model_summary import ModelSummary
|
|
from lightning.pytorch.callbacks.early_stopping import EarlyStopping
|
|
from lightning.pytorch.loggers import WandbLogger
|
|
from lightning.pytorch.callbacks import ModelCheckpoint
|
|
from pytorch_forecasting.data.timeseries import TimeSeriesDataSet
|
|
from pytorch_forecasting import QuantileLoss
|
|
|
|
from ml.loss import GMADL
|
|
from ml.model import get_model
|
|
|
|
|
|
def get_args():
|
|
parser = argparse.ArgumentParser(
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
|
|
parser.add_argument(
|
|
"config", help="Experiment configuration file in yaml format.")
|
|
|
|
parser.add_argument(
|
|
"-p",
|
|
"--project",
|
|
default="wne-masters-thesis-testing",
|
|
help="W&B project name.")
|
|
|
|
parser.add_argument(
|
|
"-l",
|
|
"--log-level",
|
|
default=logging.INFO,
|
|
type=int,
|
|
help="Sets the log level.")
|
|
|
|
parser.add_argument(
|
|
"-s",
|
|
"--seed",
|
|
default=42,
|
|
type=int,
|
|
help="Random seed for the training.")
|
|
|
|
parser.add_argument(
|
|
"-n",
|
|
"--log-interval",
|
|
default=100,
|
|
type=int,
|
|
help="Log every n steps."
|
|
)
|
|
|
|
parser.add_argument(
|
|
'-v',
|
|
'--val-check-interval',
|
|
default=300,
|
|
type=int,
|
|
help="Run validation every n batches."
|
|
)
|
|
|
|
parser.add_argument(
|
|
'-t',
|
|
'--patience',
|
|
default=5,
|
|
type=int,
|
|
help="Patience for early stopping."
|
|
)
|
|
|
|
parser.add_argument('--no-wandb', action='store_true',
|
|
help='Disables wandb, for testing.')
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def get_dataset(config, project):
|
|
artifact_name = f"{project}/{config['data']['dataset']}"
|
|
artifact = wandb.Api().artifact(artifact_name)
|
|
base_path = artifact.download()
|
|
logging.info(f"Artifacts downloaded to {base_path}")
|
|
|
|
name = artifact.metadata['name']
|
|
part_name = f"in-sample-{config['data']['sliding_window']}"
|
|
data = pd.read_csv(os.path.join(
|
|
base_path, name + '-' + part_name + '.csv'))
|
|
logging.info(f"Using part: {part_name}")
|
|
|
|
# TODO: Fix in dataset
|
|
data['weekday'] = data['weekday'].astype('str')
|
|
data['hour'] = data['hour'].astype('str')
|
|
|
|
validation_part = config['data']['validation']
|
|
logging.info(f"Using {validation_part} of in sample part for validation.")
|
|
train_data = data.iloc[:int(len(data) * (1 - validation_part))]
|
|
val_data = data.iloc[len(train_data) - config['past_window']:]
|
|
logging.info(f"Trainin part size: {len(train_data)}")
|
|
logging.info(
|
|
f"Validation part size: {len(val_data)} "
|
|
+ f"({len(data) - len(train_data)} + {config['past_window']})")
|
|
|
|
logging.info("Building time series dataset for training.")
|
|
train = TimeSeriesDataSet(
|
|
train_data,
|
|
time_idx=config['data']['fields']['time_index'],
|
|
target=config['data']['fields']['target'],
|
|
group_ids=config['data']['fields']['group_ids'],
|
|
min_encoder_length=config['past_window'],
|
|
max_encoder_length=config['past_window'],
|
|
min_prediction_length=config['future_window'],
|
|
max_prediction_length=config['future_window'],
|
|
static_reals=config['data']['fields']['static_real'],
|
|
static_categoricals=config['data']['fields']['static_cat'],
|
|
time_varying_known_reals=config['data']['fields'][
|
|
'dynamic_known_real'],
|
|
time_varying_known_categoricals=config['data']['fields'][
|
|
'dynamic_known_cat'],
|
|
time_varying_unknown_reals=config['data']['fields'][
|
|
'dynamic_unknown_real'],
|
|
time_varying_unknown_categoricals=config['data']['fields'][
|
|
'dynamic_unknown_cat'],
|
|
randomize_length=False
|
|
)
|
|
|
|
logging.info("Building time series dataset for validation.")
|
|
val = TimeSeriesDataSet.from_dataset(
|
|
train, val_data, stop_randomization=True)
|
|
|
|
return train, val
|
|
|
|
|
|
def get_loss(config):
|
|
loss_name = config['loss']['name']
|
|
|
|
if loss_name == 'Quantile':
|
|
return QuantileLoss(config['loss']['quantiles'])
|
|
|
|
if loss_name == 'GMADL':
|
|
return GMADL(
|
|
a=config['loss']['a'],
|
|
b=config['loss']['b']
|
|
)
|
|
|
|
raise ValueError("Unknown loss")
|
|
|
|
|
|
def main():
|
|
args = get_args()
|
|
logging.basicConfig(level=args.log_level)
|
|
pl.seed_everything(args.seed)
|
|
|
|
run = wandb.init(
|
|
project=args.project,
|
|
config=args.config,
|
|
job_type="train",
|
|
mode="disabled" if args.no_wandb else "online"
|
|
)
|
|
config = run.config
|
|
logging.info("Using experiment config:\n%s", pprint.pformat(config))
|
|
|
|
# Get time series dataset
|
|
train, valid = get_dataset(config, args.project)
|
|
logging.info("Train dataset parameters:\n" +
|
|
f"{pprint.pformat(train.get_parameters())}")
|
|
|
|
# Get loss
|
|
loss = get_loss(config)
|
|
logging.info(f"Using loss {loss}")
|
|
|
|
# Get model
|
|
model = get_model(config, train, loss)
|
|
logging.info(f"Using model {config['model']['name']}")
|
|
logging.info(f"{ModelSummary(model)}")
|
|
logging.info(
|
|
"Model hyperparameters:\n" +
|
|
f"{pprint.pformat(model.hparams)}")
|
|
|
|
# Checkpoint for saving the model
|
|
checkpoint_callback = ModelCheckpoint(
|
|
monitor='val_loss',
|
|
save_top_k=3,
|
|
mode='min',
|
|
)
|
|
|
|
# Logger for W&B
|
|
wandb_logger = WandbLogger(
|
|
project=args.project,
|
|
experiment=run,
|
|
log_model="all") if not args.no_wandb else None
|
|
|
|
early_stopping = EarlyStopping(
|
|
monitor="val_loss",
|
|
mode="min",
|
|
patience=args.patience)
|
|
|
|
batch_size = config['batch_size']
|
|
logging.info(f"Training batch size {batch_size}.")
|
|
|
|
epochs = config['max_epochs']
|
|
logging.info(f"Training for {epochs} epochs.")
|
|
|
|
trainer = pl.Trainer(
|
|
accelerator="auto",
|
|
max_epochs=epochs,
|
|
logger=wandb_logger,
|
|
callbacks=[
|
|
checkpoint_callback,
|
|
early_stopping
|
|
],
|
|
log_every_n_steps=args.log_interval,
|
|
val_check_interval=args.val_check_interval
|
|
)
|
|
|
|
if epochs > 0:
|
|
logging.info("Starting training:")
|
|
trainer.fit(
|
|
model,
|
|
train_dataloaders=train.to_dataloader(
|
|
batch_size=batch_size,
|
|
num_workers=3,
|
|
),
|
|
val_dataloaders=valid.to_dataloader(
|
|
batch_size=batch_size, train=False, num_workers=3
|
|
))
|
|
|
|
# Run validation with best model to log min val_loss
|
|
# TODO: Maybe use different metric like min_val_loss ?
|
|
ckpt_path = trainer.checkpoint_callback.best_model_path or None
|
|
trainer.validate(model, dataloaders=valid.to_dataloader(
|
|
batch_size=batch_size, train=False, num_workers=3),
|
|
ckpt_path=ckpt_path)
|
|
|
|
# Clean up models that do not have best/latest tags, to save space on wandb
|
|
for artifact in wandb.Api().run(run.path).logged_artifacts():
|
|
if artifact.type == "model" and not artifact.aliases:
|
|
logging.info(f"Deleting artifact {artifact.name}")
|
|
artifact.delete()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|