amazon-science / unconditional-time-series-diffusion

Official PyTorch implementation of TSDiff models presented in the NeurIPS 2023 paper "Predict, Refine, Synthesize: Self-Guiding Diffusion Models for Probabilistic Time Series Forecasting"
Apache License 2.0
101 stars 21 forks source link

Generate the multivariate sequences #11

Open hanli997 opened 1 month ago

hanli997 commented 1 month ago

Thank you for your excellent work and for sharing the codebase. I have two questions regarding the paper: 1.I am interested in generating multivariate sequences. In the paper, you mentioned that “TSDiff can be modified to handle multivariate time series by incorporating additional layers, e.g., a transformer layer, operating across the feature dimensions after the S4 layer.” Could you please provide further explanation on how to implement this? 2.Is the format of the synthetic data the same as the original data? For instance, in TimeGAN, the original sequences are divided into multiple idd segments, and the synthetic data generated is also in multiple segments, rather than a complete time series. Please excuse any ignorance on my part; I am eager to learn more about your work.

hanlaoshi commented 1 month ago

Hi, tsdiff is an outstanding work, and I have also been studying this paper recently. I am also very interested in generating multivariate sequences with tsdiff. Based on my experience with GluonTS, I tried implementing a multivariate sequence. The training process went smoothly, but I encountered a mismatch in shapes between data['future_target'] and scaled during the evaluation stage. Here is the code from my train_model.py. You will also need to modify the input_dim and output_dim in the configs.py file, as well as the create_transforms function in utils.py regarding FieldName.TARGET's expected_ndim, changing expected_ndim=1 to expected_ndim=2. We can discuss any issues together. I am also consulting the authors about the shape mismatch issue during testing and look forward to their response.

train_model.py

def main(config, log_dir):
    # Load parameters
    dataset_name = config["dataset"]
    freq = config["freq"]
    context_length = config["context_length"]
    prediction_length = config["prediction_length"]
    total_length = context_length + prediction_length

    # Create model
    model = create_model(config)

    # Setup dataset and data loading
    # dataset = get_gts_dataset(dataset_name)
    dataset = get_dataset(dataset_name)
    train_grouper = MultivariateGrouper(max_target_dim=min(2000, int(dataset.metadata.feat_static_cat[0].cardinality)))
    test_grouper = MultivariateGrouper(num_test_dates=int(len(dataset.test)/len(dataset.train)), 
                                   max_target_dim=min(2000, int(dataset.metadata.feat_static_cat[0].cardinality)))
    train_dataset = dataset.train
    test_dataset = dataset.test
    dataset_train = train_grouper(train_dataset)
    dataset_test = test_grouper(test_dataset)# type: ignore
    input_size = int(dataset.metadata.feat_static_cat[0].cardinality)
    assert dataset.metadata.freq == freq
    assert dataset.metadata.prediction_length == prediction_length

    if config["setup"] == "forecasting":
        training_data = dataset_train
    elif config["setup"] == "missing_values":
        missing_values_splitter = OffsetSplitter(offset=-total_length)
        training_data, _ = missing_values_splitter.split(dataset_train)

    num_rolling_evals = int(len(dataset.test) / len(dataset.train))

    transformation = create_transforms(
        num_feat_dynamic_real=0,
        num_feat_static_cat=0,
        num_feat_static_real=0,
        time_features=model.time_features,
        prediction_length=config["prediction_length"],
    )

    training_splitter = create_splitter(
        past_length=config["context_length"] + max(model.lags_seq),
        future_length=config["prediction_length"],
        mode="train",
    )

    callbacks = []
    if config["use_validation_set"]:
        transformed_data = transformation.apply(training_data, is_train=True)
        # Assuming `training_data` is a Map or similar iterable object
        for entry in training_data:
            print("Shape of target data before transformation:", entry['target'].shape)
            break  # We only want to check the first entry for debugging
        # Assuming the transformation outputs another iterable dataset
        # for entry in transformed_data:
        #     print("Shape of target data after transformation:", entry['target'].shape)
        #     break  # Again, just check the first entry
        # print("Shape of target data before transformation:", training_data[0]['target'].shape)
        # print("Shape of transformed data:", transformed_data[0]['target'].shape)
        train_val_splitter = OffsetSplitter(
            offset=-config["prediction_length"] * num_rolling_evals
        )
        _, val_gen = train_val_splitter.split(training_data)
        val_data = val_gen.generate_instances(
            config["prediction_length"], num_rolling_evals
        )

        callbacks = [
            EvaluateCallback(
                context_length=config["context_length"],
                prediction_length=config["prediction_length"],
                sampler=config["sampler"],
                sampler_kwargs=config["sampler_params"],
                num_samples=config["num_samples"],
                model=model,
                transformation=transformation,
                test_dataset=dataset_test,
                val_dataset=val_data,
                eval_every=config["eval_every"],
            )
        ]
    else:
        transformed_data = transformation.apply(training_data, is_train=True)

    log_monitor = "train_loss"
    filename = dataset_name + "-{epoch:03d}-{train_loss:.3f}"

    data_loader = TrainDataLoader(
        Cached(transformed_data),
        batch_size=config["batch_size"],
        stack_fn=batchify,
        transform=training_splitter,
        num_batches_per_epoch=config["num_batches_per_epoch"],
    )

    checkpoint_callback = ModelCheckpoint(
        save_top_k=1,
        monitor=f"{log_monitor}",
        mode="min",
        filename=filename,
        save_last=True,
        save_weights_only=True,
    )

    callbacks.append(checkpoint_callback)
    callbacks.append(RichProgressBar())

    trainer = pl.Trainer(
        accelerator="gpu" if torch.cuda.is_available() else None,
        devices=[int(config["device"].split(":")[-1])],
        max_epochs=config["max_epochs"],
        enable_progress_bar=True,
        num_sanity_val_steps=0,
        callbacks=callbacks,
        default_root_dir=log_dir,
        gradient_clip_val=config.get("gradient_clip_val", None),
    )
    logger.info(f"Logging to {trainer.logger.log_dir}")
    trainer.fit(model, train_dataloaders=data_loader)
    logger.info("Training completed.")

    best_ckpt_path = Path(trainer.logger.log_dir) / "best_checkpoint.ckpt"

    if not best_ckpt_path.exists():
        torch.save(
            torch.load(checkpoint_callback.best_model_path)["state_dict"],
            best_ckpt_path,
        )
    logger.info(f"Loading {best_ckpt_path}.")
    best_state_dict = torch.load(best_ckpt_path)
    model.load_state_dict(best_state_dict, strict=True)

    metrics = (
        evaluate_guidance(config, model, dataset_test, transformation)
        if config.get("do_final_eval", True)
        else "Final eval not performed"
    )
    with open(Path(trainer.logger.log_dir) / "results.yaml", "w") as fp:
        yaml.dump(
            {
                "config": config,
                "version": trainer.logger.version,
                "metrics": metrics,
            },
            fp,
        )
abdulfatir commented 4 weeks ago

Hi @hanlaoshi! Apologies for the very late reply. The current codebase does not support multivariate training and evaluation out of the box. You would have to make modifications to enable multivariate. If you can share a small working example with the issue you are facing, we could take a look at it to resolve the issue but do note that multivariate time series are beyond the scope of this code.

cc @marcelkollovieh

abdulfatir commented 4 weeks ago

@hanli997

  1. It may not be a straightforward modification in the current codebase. The guidance and refinement scheme are general and should work with multivariate series. You'd need to modify the architecture to mix information from different dimensions.
  2. TSDiff will generate sequences of a fixed length.
hanlaoshi commented 4 weeks ago

Hi @hanlaoshi! Apologies for the very late reply. The current codebase does not support multivariate training and evaluation out of the box. You would have to make modifications to enable multivariate. If you can share a small working example with the issue you are facing, we could take a look at it to resolve the issue but do note that multivariate time series are beyond the scope of this code.

cc @marcelkollovieh

Hi, thank you very much for your response. As guided in the paper, additional layers running across feature dimensions can be added to accommodate multivariate time series. We have currently modified TSDiff to handle multivariate scenarios, but we have found that its performance may be significantly weaker than task-specific models like CSDI in these scenarios. Could you help us analyze the possible reasons for this?

hanli997 commented 4 weeks ago

@abdulfatir Thank you for your response. I will make the necessary attempts based on your suggestions. Looking forward to your continued exceptional work.