Merge pull request #7 from JonathanCrabbe/other_backbone

Other Backbone

cmd/conf/score_model/lstm.yaml

@@ -0,0 +1,10 @@
+_target_: fdiff.models.score_models.LSTMScoreModule
+_partial_: true
+d_model: 72
+num_layers: 10
+lr_max: 1.0e-3
+fourier_noise_scaling: ${fourier_transform}
+likelihood_weighting: False
+
+defaults:
+  - noise_scheduler: vpsde

cmd/conf/score_model/mlp.yaml

@@ -0,0 +1,11 @@
+_target_: fdiff.models.score_models.MLPScoreModule
+_partial_: true
+d_model: 72
+d_mlp: 1024
+num_layers: 10
+lr_max: 1.0e-4
+fourier_noise_scaling: ${fourier_transform}
+likelihood_weighting: False
+
+defaults:
+  - noise_scheduler: vpsde

cmd/sample.py

@@ -11,7 +11,7 @@
 from fdiff.models.score_models import ScoreModule
 from fdiff.sampling.metrics import MetricCollection
 from fdiff.sampling.sampler import DiffusionSampler
-from fdiff.utils.extraction import dict_to_str, get_best_checkpoint
+from fdiff.utils.extraction import dict_to_str, get_best_checkpoint, get_model_type
 from fdiff.utils.fourier import idft
 
 
@@ -48,7 +48,8 @@ def __init__(self, cfg: DictConfig) -> None:
 
         # Load score model from checkpoint
         best_checkpoint_path = get_best_checkpoint(self.save_dir / "checkpoints")
-        self.score_model = ScoreModule.load_from_checkpoint(
+        model_type = get_model_type(train_cfg)
+        self.score_model = model_type.load_from_checkpoint(
             checkpoint_path=best_checkpoint_path
         )
         if torch.cuda.is_available():