Implement simple model checkpointing

[meffmadd]

No description provided.

[nmichlo]

Hi @meffmadd, thank you so much for this contribution. It looks great! 😁🎉

Before we merge, please may you update the checkpointing in the configs as some tests are currently failing and for documentation purposes?

• Both in the main config (possibly defaulting to false to maintain current behavior)

  • disent/experiment/config/config.yaml

    Lines 44 to 47 in 8dee583

    	# TODO: https://pytorch-lightning.readthedocs.io/en/stable/common/weights_loading.html
    	# checkpoint:
    	# load_checkpoint: NULL # NULL or string
    	# save_checkpoint: FALSE # boolean, save at end of run -- more advanced checkpointing can be done with a callback!

• And copying this to the test config in the corresponding location (defaulting to true so we know the checkpoint behavior works).
  • https://github.com/nmichlo/disent/blob/8dee583993d582729928984640930a4e92e65de6/experiment/config/config_test.yaml

On that note, we might want to just add the checkpointing to the pytorch lightning train step, using the same hook. Then possibly at the end, load the checkpoint after everything to make sure it works.

• disent/tests/test_frameworks.py

  Lines 98 to 126 in 8dee583

  	@pytest.mark.parametrize(['Framework', 'cfg_kwargs', 'Data'], _TEST_FRAMEWORKS)
  	def test_frameworks(Framework, cfg_kwargs, Data):
  	DataSampler = {
  	1: GroundTruthSingleSampler,
  	2: GroundTruthPairSampler,
  	3: GroundTruthTripleSampler,
  	}[Framework.REQUIRED_OBS]
  	data = XYObjectData() if (Data is None) else Data()
  	dataset = DisentDataset(data, DataSampler(), transform=ToImgTensorF32())
  	dataloader = DataLoader(dataset=dataset, batch_size=4, shuffle=True, num_workers=0)
  	framework = Framework(
  	model=AutoEncoder(
  	encoder=EncoderLinear(x_shape=data.x_shape, z_size=6, z_multiplier=2 if issubclass(Framework, Vae) else 1),
  	decoder=DecoderLinear(x_shape=data.x_shape, z_size=6),
  	),
  	cfg=Framework.cfg(**cfg_kwargs)
  	)
  	# test pickling before training
  	pickle.dumps(framework)
  	# train!
  	trainer = pl.Trainer(logger=False, checkpoint_callback=False, max_steps=256, fast_dev_run=True)
  	trainer.fit(framework, dataloader)
  	# test pickling after training, something may have changed!
  	pickle.dumps(framework)

[meffmadd]

Hi, I couldn't work on it this weekend but will start now and fix the configs so that the tests work again. I will also add a test case that tests the behavior.

[codecov]

Codecov Report

Base: 70.01% // Head: 70.04% // Increases project coverage by +0.02% 🎉

Coverage data is based on head (28d0e95) compared to base (8dee583).
Patch coverage: 100.00% of modified lines in pull request are covered.

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #37      +/-   ##
==========================================
+ Coverage   70.01%   70.04%   +0.02%     
==========================================
  Files         135      135              
  Lines        7531     7538       +7     
==========================================
+ Hits         5273     5280       +7     
  Misses       2258     2258              

Impacted Files	Coverage Δ
disent/frameworks/_ae_mixin.py	90.00% <100.00%> (+1.32%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[nmichlo]

Great work! Thank you so much for making these changes!