Factorize Models to Prepare HF Upload

[LTMeyer]

Context

We want to upload the checkpoints of the model trained for the benchmark to HF (see #33).
To ease the loading of the checkpoints from_pretrained(<hf_url>) and not having to pass any redundant parameters this requires all the parameters of the model class to be of standard type. This is not the case as we use the dataset.metadata for instantiation. This PR fixes it by setting explicitly the n_spatial_dims and spatial_resolution arguments that were originally extracted from the metadata.

Content

• Add a BaseModel class to support HF upload;
• Change dset_metadata -> actual n_spatail_dims and spatial_resolution arguments:
• Add tests for most of the models.

[francois-rozet]

The approach seems sound to me, if the training scripts still run. I just left a minor comment.

[LTMeyer]

The approach seems sound to me, if the training scripts still run.

I've just tried the following and it worked as expected.
python train.py experiment=fno server=rusty data=active_matter trainer=debug

[rubenohana]

To be sure, does it work for 3D data as well?

[mikemccabe210]

Looks like a solid improvement to me, especially the test additions!

[LTMeyer]

To be sure, does it work for 3D data as well?

It should. If you have a command you would like me to test let me know.

[rubenohana]

To be sure, does it work for 3D data as well?

It should. If you have a command you would like me to test let me know.

Could you try the same command for training on the MHD_64 dataset?

[LTMeyer]

Looks like a solid improvement to me, especially the test additions!

With the tests I've noticed by the way that some models seem to accept B H W C while other expect B C H W as input shapes (see https://github.com/PolymathicAI/the_well/pull/34/files#diff-ca99d4b52e8a6a4e8fdadd6871c1152341a654ab0c18f503fbdb715820cad05fR58 and https://github.com/PolymathicAI/the_well/pull/34/files#diff-e6fd5adf396bade4ea5c4661206b237722692d3f31beb865ed0720ea149a76e4R17-R21). I don't remember if this was expected?

[LTMeyer]

To be sure, does it work for 3D data as well?

It should. If you have a command you would like me to test let me know.

Could you try the same command for training on the MHD_64 dataset?

I've just checked the following command runs as expected python the_well/benchmark/train.py experiment=fno server=rusty data=MHD_64 trainer=debug.