[STEP] load and save in BasePkg in pytorch-forecasting v2

[phoeenniixx]

This EP proposes a design for loading and saving methods in BasePkg in pytorch-forecasting v2 - to load and save models and other artifacts (like scalers)

Based on hackmd: https://hackmd.io/@Pm5-sJBvSfeR6I59oCaLOA/Bk9Xt5Txzl/edit

[phoeenniixx]

Moving the converstation from the hackmd here:
@agobbifbk commented:
BasePkg is the model layer right? So in the save method we need to extract the scalers from the data (D2 layer instance). This is the coupling that worries Franz right? What if we just save the checkpoints at BasePkg layer and let the D2 save its own artifacts? If you are worried about discrepancy between paths we can think on an Experiment layer that requires a D1 a D2 and a model layer classes where a single path is needed and shared between layers eventually. Otherwise you are saving D2 stuff in the model layer and this can cause a lot of trouble when we add other D2 layers no?

@fkiraly replied:
yes, what the D2 layer saves will depend on the particular D2 class. So, the pkg layer cannot contain logic for all D2 classes. Instead, there needs to be an abstract method in D2 that saves its artefacts, and the pkg layer can call it using the strategy pattern.

@phoeenniixx replied:

BasePkg is the model layer right?

Not quite! This is a wrapper around the model layer that provides interface to add tests and fit and predict without actually touching the model layer. See complete details here: #46

So in the save method we need to extract the scalers from the data (D2 layer instance). This is the coupling that worries Franz right?
yes, what the D2 layer saves will depend on the particular D2 class. So, the pkg layer cannot contain logic for all D2 classes. Instead, there needs to be an abstract method in D2 that saves its artefacts, and the pkg layer can call it using the strategy pattern.

Yes, the idea to create a rather generic method that can handle different D2 modules and different models as well.
So, I think a tree of inheritance is the best way to handle such kind of issues, we would have the most generic parts higher in the heirarchy and we specialise as we move down

[fkiraly]

does this mean a checkpoint is always saved, i.e., a user has no means to turn off checkpoint saving? I was assuming the default is "it is not saved". Otherwise, there is silently a file written to the hard drive on every fit call with defaults?

[phoeenniixx]

no, you have to pass save_ckpt=True, it is False by default. It will be saved ONLY IF you pass this arg as True.

[agobbifbk]

Just one comment, but I don't want to overcomplicated the codebase: as I can imagine you are saving the checkpoint that leads to the lowest validation loss (as far as I remember it is related to the trainer), but in some cases, this minimum can be an 'false' minumum and what you would like to save is also the last checkpoint. I know that this would double the space of the saved weights, but sometimes you would let the network continue till the end. It is just a comment, not a mandatory change

[agobbifbk]

Ah I see at row 406 that I can eventually save the last checkpoint :-)

[fkiraly]

can you be more explicit on how many files and which file names this is saved to?

That would also be crucial in the docstring. A docstring should document all assumptions and guarantees, that is all preconditions that need to be satisfied, and all changes at the point of the end state it makes.

If file access or file saving is involved, then, very importantly:

• assumptions may include assumptions on files being present in a specific format or location, and the content of these files being something.
• guarantees may include files being written or changed, at specific file locations with a specific content.

[fkiraly]

what is the reason to save configs, or to assume saved configs? For inference, at least the trainer_cfg seems like it is not needed. Similarly, could a different data_cfg not also make sense?

Imo we need to think carefully about two cases:

• save and load everything
• save and load only the weights
• any intermediate option that may make sense

[phoeenniixx]

Well for data modules, we need to save the cfgs to match the exact conditions the model was trained on (like same max_encoder_length, batch_size etc). Obv the user could use their own cfgs still if they want, if they pass their own cfgs that would be given higher preference and would override the saved cfgs. But saving provides the users to not to pass the cfgs they used previously again.

About the trainer_cfg, I am a bit conflicted. I think we should add the way to save it as this would save the user from passing the cfg again if they are planning to use the same trainer they used earlier. And way to overrride it would give them complete flexibility. Although I agree saving the trainer_cfg is not something that is very important and would cause issues, but if we provide this, it would save some users (who are lazy like me) to pass the same cfg again. THey can always override it to use a new configuration for the trainer anytime.

But I think saving the cfgs for data modules are more important - it provides safeguards from the "accidents" where the user might not remember (or know) the previous cfgs for the data module and that may lead to issues

[fkiraly]

sure, still I think saving and loading some or all should be optional (where it makes sense)

[fkiraly]

in an enhancement proposal, there should be a preamble. What problem are we solving, what requirements there are, what the design needs to do. Optimally, also a conceptual model and the use cases we want to solve. E.g., what are we saving? What are we loading? And why? Are there multiple use cases related to saving and loading? I think pure inference use as well as loading in order to continue training, etc, should be considered.

[fkiraly]

question: is there a case to be made for saving to be separate from fitting? E.g., is there a use case where saving does not immediately follow fitting?

[phoeenniixx]

Tbh, I cant think of any case as such. I mean we would like to save the model only after it has learned something no? why would anyone would like to save random weights?

Can you please rephrase the question, I think I am not able to understand the exact question here

[fkiraly]

the default is "save in checkpoints dir". However, save_ckpt and ckpt_dir feel redundant, as ckpt_dir=None could indicate that no saving is desired. This would make save_ckpt unnecessary.

[phoeenniixx]

I agree, but having the save_ckpt, imo, provides a clear "feedback" to the user that only this arg is going to make the saving of models. ckpt_dir is kinda optional - only if the user wants to save the ckpts and other artifacts to somewhere specific. Usually all the checkpoints would always save in the root within a new folder- checkpoints. I dont think this ckpt_dir should decide if the user wants to save the skpts or not, for me, I would usually just pass True and would not pass anything to the ckpt_dir. ckpt_dir is just to provide more flexibility to the user.

But again this adds a little bit of redundancy. This is a question of redundancy and of being discreet. It is upto us which path we need to take.

[phoeenniixx]

If you have more inclination towards the removal of redundancy I am happy to remove it completely

[fkiraly]

ckpt_dir is kinda optional - only if the user wants to save the ckpts and other artifacts to somewhere specific.

That makes no sense to me. I can think of no scenario in which the user knows they want to save, but want no control about where it gets saved.

I still would not separate the two parameters like this.

[fkiraly]

having arguments foo and foo_kwargs always indicates that maybe we should consider using a dataclass constructor pattern instead, and pass only foo, but it can be a newly constructed object, in this case ModelCheckpoint. Why not ask users to pass a ModelCheckpoint if they want kwargs?

[phoeenniixx]

Sorry I didnt understand the question
We are not asking for foo, just foo_kwargs, we always create foo internally ourselves.
If we ask form the user to pass the ModelCheckpoint, they would themselves have to create an instance of ModelCheckpoint and I think that goes against the idea of encapsulation. If we know we are always going to use ModelCheckpoint, why ask from the user to pass it? Just ask for any specific kwargs they want to pass - which is completely optional.

This also saves the users from malicious code. If we ask the user to pass the instance of ModelCheckpoint, someone can create their own child of this class with some malicious code and pass that to the method. And as the footprint of the child is similar to ModelCheckpoint, it could cause issues - although this might be bit of a stretch, but it is possible. If the user wants so much of customization, they would bypass pkg class and use the "bare bones" pipeline instead.

I think we should always just ask for the cfgs and keep most of the things encapsulated as much as we can - to save the effort of the user to create everything themselves and also from security PoV it is better.

What do you think? @fkiraly

[fkiraly]

hm, ok, I think that makes sense. Are there no other classes the ModelCheckpoint could be replaced by?

[fkiraly]

should there be a loose function load(ckpt_path:str)? Which recognizes what is in the path, and then loads it?

Currently, the user needs to know which model is stored in a file location, which feels cumbersome, as in the case of many models they need to keep track which model is where. Instead, a bare load would, or could, be more comfortable.

[phoeenniixx]

Yeah that is the issue of the current implementation, please see Proposed design sec for more info, I have added similar method there already

[fkiraly]

I think you need to think about use cases clearly. There may be distinctions as in:

• how much the user wants to save or load - e.g., weights only, configs too, and even fitted scalers etc.
• when the user wants to save or load - e.g., to checkpoint their own training process, or simply to load an inference ready checkpoint

[phoeenniixx]

I think the tree inheritance takes care of this, but I think we should add some args for saving scalers, thanks I will make the updates

[phoeenniixx]

Hi @fkiraly, @agobbifbk I think the design is ready for review

[agobbifbk]

Are you sure that empty is the default? In terms of usability I would expect a list of all the possible things to save: if i use the default values I can not reuse the models, this feels strange to me

[phoeenniixx]

by exclude being empty, we would have an empty list - meaning NOTHING is excluded, and we save EVERYTHING.
Do you prefer it in some other way?

[agobbifbk]

My bad, it's fine like now!

[fkiraly]

I am thinking about having an exclude or an include parameter, I suppose exclude is better? But I am not sure, often the positive specification makes sense if it is an action. Default (e.g., None) saves all, and it might be easier for a user to think about what to save as opposed to what not to save, in difference to a hard to remember full set of options?

[agobbifbk]

typo here:

       # scalers are already None by default, so we dont need to pass scalers=None

[agobbifbk]

This feels strange to me, I used scalers but I exclude them, so how can I suppose to run the model in inference if I lost all the information about the scalers?

[phoeenniixx]

This is the case where the user just want to save the model weights. Nothing else, maybe useful in something like pretraining? I dont think it is always the case that we save (and use) the same scalers everytime (like in pretrianing)?

[agobbifbk]

I see, but the weight will work only if the data are scaled as the training dataset. We can keep this functionality for sure, but I would like to see a giant warning :-)

[fkiraly]

well, I think it is the user who is giving instructions explicitly in this case.

Perhaps we want some better example in the vignette on including or excluding artefacts?

[agobbifbk]

Hi @agobbifbk, I wanted to know your preference as well here, should we have a artifacts.json or do you have a better approach in mind for load?

Seems clean to me, in DSIPTS (or in other projects) we save all the configuration in a yaml file because we build our code leveraging on HYDRA https://hydra.cc/docs/intro/. The benefit of using hydra is that when you use the code via bash you can overwrite some of the configuration directly by command line. I suggest to have a look to it and see if this is an overkill for this particular task. I try to give here an example how the code reads if all the configurations are manage through omegaconf-hydra:

NOW:

    dataset = TimeSeries(
        data=data_df, # data_df is any arbitrary dataframe
        time="time_idx",
        target="y",
        group=["series_id"],
        num=["x", "future_known_feature", "static_feature"],
        cat=["category", "static_feature_cat"],
        known=["future_known_feature"],
        unknown=["x", "category"],
        static=["static_feature", "static_feature_cat"],
    )

    # data module config
    datamodule_cfg = dict(
        max_encoder_length=30,
        max_prediction_length=1,
        batch_size=32,
    )

    # model config
    model_cfg = dict(
        loss=MAE(),
        logging_metrics=[MAE(), SMAPE()],
        optimizer="adam",
        optimizer_params={"lr": 1e-3},
        lr_scheduler="reduce_lr_on_plateau",
        lr_scheduler_params={"mode": "min", "factor": 0.1, "patience": 10},
        hidden_size=64,
        num_layers=2,
        attention_head_size=4,
        dropout=0.1,
    )

    # trainer config
    trainer_cfg = dict(
        max_epochs=5,
        accelerator="auto",
        devices=1,
        enable_progress_bar=True,
        log_every_n_steps=10,
    )

    model_pkg = TFT_pkg_v2(
        model_cfg=model_cfg,
        trainer_cfg=trainer_cfg,
        datamodule_cfg=datamodule_cfg,
    )

The new schema will be something like:

    dataset = TimeSeries(
        data=data_df, # data_df is any arbitrary dataframe
        **config.dataset
      
    )
    model_selected = eval(config.model)
    model_pkg = model_selected(
        model_cfg=config.model_cfg,
        trainer_cfg=config.trainer_cfg,
        datamodule_cfg=config.datamodule_cfg,
    )

with a yaml file

paths:
          root: checkpoints ## this is the main path dir

 artifacts: 
           best_model: None ## these will be filled during the training
           scalers: None
           target_normalizers: None
           model_cfg: None
           datamodule_cfg: None
           trainer_cfg: None
           datamodule_metadata: None
                 
            
load: 
      skip: [scalers] ## use after in the when

save: 
      opt1: other_options

      

dataset:
          time: time_idx
          target: y
          group: [series_id]
          num: [x, future_known_feature, static_feature]
          cat: [category, static_feature_cat]
          known: [future_known_feature]
          unknown: [x, category]
          static: [static_feature, static_feature_cat]

data_module:
          max_encoder_length: 30
          max_prediction_length: 1
          batch_size: 32

model_cfg:
          loss: : "MAE()" # somewhere after you need to call `eval `
          logging_metrics: ["MAE()", "SMAPE()"]
          optimizer: "adam"
          optimizer_params:
                     lr: 1e-3
          lr_scheduler: "reduce_lr_on_plateau"
          lr_scheduler_params: 
                   mode: "min"
                   factor: 0.1
                   patience: 10
          ## these are somehow model dependent, hydra can let you manage this, meaning that you can have, for each model, a dedicated config file with its own parameter and than hydra manage to glue together the different configs
          hidden_size: 64
          num_layers: 2
          attention_head_size: 4
          dropout: 0.1
model: TFT_pkg_v2
          

 trainer_cfg:
          max_epochs: 5
          accelerator: auto
          devices: 1
          enable_progress_bar: True
          log_every_n_steps: 10
    

The main idea is that the artifact section will be filled during the training and afterwards you just need to save the config. In this way the user can see HOW he trained the model, and reload it afterwards pointing to the saved config file, eventually overwriting some parameters (e.g. load.skip=[scalers]). Some colleagues have different files for the different keys of the configuration, this can be useful for a cleaner structure, here I glue all the config together.

This may seems an overkill but please have a look to this: https://hydra.cc/docs/tutorials/basic/running_your_app/multi-run/ in my opinion this is a nice-to-have feature!

[agobbifbk]

Another option is: close this with artifacts.json and then open another issue for managing the configurations :-) it seems the most reasonable path

[phoeenniixx]

I think hydra is a nice-to-have feat, but a feat that should be a "future" endeavour - once we have atleast the basic API stable. But I think artifacts could be a yaml in place of json? Is there any "preferred" file format , that can make things easier for the model deployments?

[agobbifbk]

I think that json or yaml are equivalent at this stage but yaml opens better compatibility in case we will switch to hydra in the future. Moreover, the user potentially will access to the V2 api through a configuration mechanism, hopefully not args from the main, so yaml is definitively the most reasonable choice here IMHO

[phoeenniixx]

Thanks, I think this also resonates with how the configs (model_config, datamodule_cfg etc) can be used - we can easily pass yaml files that contain these configs to the methods and it would work just fine. So, artifacts could be a yaml in place of json and it would keep things uniform across the pipeline.

Thanks I will update the proposal with this new change!

[fkiraly]

what does "best model" mean here? Should it mot simply be model.ckpt for simplicity?

[phoeenniixx]

i took it from the "naming" convention - mainly we save the best best model with best performance, But i agree I should add a generalised name here in the example.

But I would disagree if you mean that we save the model ckpts as model.ckpt simply. As sometimes we dont save just one model - the best or last one, but we save the models at every one (or few epochs). In that case it would be useful to have a naming that clearly define the epoch, step etc

[phoeenniixx]

But I think this is more of a ModelCheckpoint responsibility then ours.

[fkiraly]

agreed!

[fkiraly]

you do not need to cite me here specifically, but you can list me as a contributor at the top.

[fkiraly]

I think json and yaml should be options in the end

[phoeenniixx]

Hmm, so you are suggesting we add support for both, that is a good idea. But rather than adding a new param to decide which file ext to use, would it make more sense to allow json and the method is json compatible but save would always give out a yaml file.

Meaning if the user wants they can pass their own json file, but if they are relying on pkg.save, it would always give out yaml file.

Is this a good idea or am I just oversimplifying things and it would be better to have a param to let the user decide - what save gives out - yaml or json?

I am not opposed to any of the solutions, I just want to know if this would be a good idea to make save a bit strict and keep its param set small?

[fkiraly]

similarly, I think json should also be ultimately supported. In load, we can disambiguate by format or file ending.

[fkiraly]

hm, I think this is a point of confusion!

I feel the behaviour here needs to be made much clearer as to what gets saved!

We may also want to have a config that determines what is the state of the python object after fit, e.g., best model or last epoch. If this does not exist, it would be very confusing that you can have behaviour after saving that you cannot have in-memory (e.g., a choice of "best model" vs "last epoch")

[phoeenniixx]

Hmm, then I think it would be best that we leave this to ModelCheckpoint? If you pass save_last=True in that callback, it would save the model from the last epoch. Actually ModelCheckpoint can help save model ckpts in any way (can save the best model as well based on different "monitor"" etc), so should we just remove pkg.save as a endpoint and let fit call it only?
See the documentation of ModelCheckpoint here: https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html

If you call save after fit, there would be only the model from the last epoch that would be present in the memory (in most cases) and that is the only thing that save could save at that time. ANd for that we would need to add specific ways as ModelCheckpoint is a callback and can save the model via Trainer, but after fit we wouldnt be using Trainer.

I think it is not worth the effort to keep save public?
Only load stays public and save becomes private called by fit.

What do you think about this?

[fkiraly]

confusing, since best_model seems to depend on context. I would say, this should be model.ckpt, since whether the model was "best" in a sense, or just "some" model, is not an intrinsic property of the artefacts saved.

[fkiraly]

• how does the user know what options there are in exclude? This should be clear.
• in particular, I would make include/exclude strings equal to the names of files, where possible.

I think this part of the design still requires some detailed thinking to make it easy to use, but we are probably not far from a good solution.

[phoeenniixx]

I think exclude is better than include as while running we would want to save everything in most cases, and added everything to a list is a waste of effort, imo. I think we should think of a rarer case where we dont save a specific thing - so we exclude it. Also, in most cases include would contain 3-4 items (model checkpoints, scalers, target_normalizers, etc) but exclude would be a smaller list.

• how does the user know what options there are in exclude? This should be clear.

It would be ANY artifact - except the configs and datamodule.metadata. (and maybe even model chekcpoints?)

• in particular, I would make include/exclude strings equal to the names of files, where possible.

I am sorry I am not very clear what you mean by this, can you please elaborate on this?

[fkiraly]

these missing indentations are bothering my mental linter...

[fkiraly]

just for structure, I would show a simple save/load cycle as a first vignette pair. Only then I would present more variants of either.

The reason for this is that save and load need to match each other.

[fkiraly]

this is weird - is this not equivalent to simply removing the files from the manifest? Why add a skip instead?

[fkiraly]

disagree. I think user choice of what to load should not require file manipulation!

[fkiraly]

if the user didn't want something persisted, they shouldn't have saved it

I think this is an incorrect argument.

You need to consider the use case where someone else has saved the artefacts. Standard use case in loading weights from huggingface!

[phoeenniixx]

Hmm, so I think you and @agobbifbk both agree with the alternative design (see line 835) of adding a param instead?
I would go forward with that then

[fkiraly]

Really great! Especially the idea of multiple files, defined folder structure, and manifest file (artifacts,yml).

I have two major questions or issues with the design:

• loading only parts of a saved file package should not require manipulating the artifacts.yml file, see discussion above
• I think the semantics vs file format of "best model" and "last model" need to be clarified - on both layers, in-memory (what is the state after fit?), and after save (what is the state after save?). Plus, whether it makes sense to distinguish "best models" from "last models" after save. I currently think it makes no sense to have a distinction of the save file format by nature of extrinsic origin.

I also left some smaller comments.

[fkiraly]

I think none of the three options here make sense to me.

• "just exclude in save" - assumes the person saving and loading are the same. Cannot assume that
• the remaining two: modification of loading behaviour which as a concern sits in the python layer should not require file modifications.

[fkiraly]

in my opinion, you already see how bad an idea the options are by how bad the (very honest and precise) docstring gets.

Noting that the docstring is "bad" since these are very unpleasant instructions, but it is also "very good" in that it captures with precision what you want the user to do. Like a very precise recipe on how to make a sandwich with canned tuna, mustard, strawberries, and woodchips.

[fkiraly]

I think you need to think about the assumption when loading the model. Having a variable name for files in checkpoints feels like a recipe for imprecise specifications. I would set this to model.ckpt and not vary the name depending on origin. If you want to add metadata on origin, this could be in a separate file (with standardized name) but I would even not do that.

[phoeenniixx]

here X and Y were just "placeholders" in reality it would be something like this: best-epoch=3-step=100.ckpt (best model from epoch 3, step 100).

But I think this is somehthing that ModelCHeckpoint should (and would) handle. we just need to pass filename as the params to it

[phoeenniixx]

Hi @fkiraly @agobbifbk
After discussion (see above) I think these are the changes and point of discussions:

• We move forward with skip param in load and no need of user to manipulate artifacts.json? [STEP] load and save in BasePkg in pytorch-forecasting v2 #51 (comment)
• Still need to decide exaclty how we support both json and yaml files for artifacts. See this comment: [STEP] load and save in BasePkg in pytorch-forecasting v2 #51 (comment)
• Still need to decide how exclude should work: [STEP] load and save in BasePkg in pytorch-forecasting v2 #51 (comment)
• Still need to decide if save should be private or public: [STEP] load and save in BasePkg in pytorch-forecasting v2 #51 (comment)

Please let me know what you think about these points

[agobbifbk]

Hi @fkiraly @agobbifbk After discussion (see above) I think these are the changes and point of discussions:

• We move forward with skip param in load and no need of user to manipulate artifacts.json? [STEP] load and save in BasePkg in pytorch-forecasting v2 #51 (comment)
• Still need to decide exaclty how we support both json and yaml files for artifacts. See this comment: [STEP] load and save in BasePkg in pytorch-forecasting v2 #51 (comment)

As mentioned in the other comments, for me yaml offers more compatibility for future development

• Still need to decide how exclude should work: [STEP] load and save in BasePkg in pytorch-forecasting v2 #51 (comment)

We can also consider to explicity pass the parameters like:

best_model = model_pkg.fit(
        dataset, # a TimeSeries Object of dataset
        ckpt_dir=ckpt_dir, # not None means we do the checkpointing
        model_ckpt_kwargs={"monitor": "train_loss_epoch"}, # for ModelCheckpoint
        save_scalers = True,
        save_target_normalizer = True,
        )

and have the same flags in the load function? I need to raise also another option here: are we sure we want to save stuff during the fit phase? See next comment

• Still need to decide if save should be private or public: [STEP] load and save in BasePkg in pytorch-forecasting v2 #51 (comment)

For me it is more intuitive that the user call save when we want to save the model-scaler-etc. Suppose it has a loop that test 100 parameters and checks for the validation loss and he wants to keep/save only one of those. If you save under the hood for each iteration you will end up with 100 .pt files... We can store temporary checkpoints in a folder if needed (eventually delete it after save is called) and when the user call save specifying a path we copy the best and last checkpoint in that folder and save scalers and normalizer eventually. But I'm not a software engineer, I don't have any strong position on this.

Please let me know what you think about these points

I hope this can help and not add noise/confusion :-)

[phoeenniixx]

For me it is more intuitive that the user call save when we want to save the model-scaler-etc. Suppose it has a loop that test 100 parameters and checks for the validation loss and he wants to keep/save only one of those. If you save under the hood for each iteration you will end up with 100 .pt files... We can store temporary checkpoints in a folder if needed (eventually delete it after save is called) and when the user call save specifying a path we copy the best and last checkpoint in that folder and save scalers and normalizer eventually. But I'm not a software engineer, I don't have any strong position on this.

Hi @agobbifbk, here we are using ModelCheckpoint (see documentation here) as a callback to trainer. I think calling save after fit would lead to the issue that we wont be able to use this callback. It would work best with fit. But I agree we need to make sure we dont oversave scalers here. And we I think we could add a flag or something to make sure if scaler is already there, we wont save it. We can even read the artifacts.yaml to make sure we dont save scalers mutliple times - if there is already scalers in yaml we wont save it again.

ModelCheckpoint allows you to save the last checkpoint, best checkpoint - everything - so I think it would be better that the user calls the save via fit and we make save private rather than public.

What do you think about this?
@agobbifbk @fkiraly

[agobbifbk]

I don't have a strong opinion on that, I always prefer to call save by myself, probably due to sklearn-catboost-xgboost paradigm :-) You still can save the path after the fit, and use them afterwards in the save, you won't lose the reference to the saved checkpoint. But it is ok to let the fit also save the stuff, are the ModelCheckpoint's params exposed to the user? What if he does not use it? Or we just expose some of them?