project Refactor, CNN Classifier Basics

lib/utils/model_io.py

@@ -1,5 +1,7 @@
 from argparse import Namespace
 from pathlib import Path
+
+import torch
 from natsort import natsorted
 from torch import nn
 
@@ -35,30 +37,25 @@ class ModelParameters(Namespace):
 class SavedLightningModels(object):
 
     @classmethod
-    def load_checkpoint(cls, models_root_path, model, n=-1, tags_file_path=''):
+    def load_checkpoint(cls, models_root_path, model=None, n=-1, tags_file_path=''):
         assert models_root_path.exists(), f'The path {models_root_path.absolute()} does not exist!'
         found_checkpoints = list(Path(models_root_path).rglob('*.ckpt'))
 
         found_checkpoints = natsorted(found_checkpoints, key=lambda y: y.name)
+        if model is None:
+            model = torch.load(models_root_path / 'model_class.obj')
+        assert model is not None
 
-        if not tags_file_path:
-            tag_files = models_root_path.rglob('meta_tags.csv')
-            tags_file_path = list(tag_files)[0]
-
-        return cls(weights=found_checkpoints[n], model=model, tags=tags_file_path)
+        return cls(weights=found_checkpoints[n], model=model)
 
     def __init__(self, **kwargs):
         self.weights: str = kwargs.get('weights', '')
-        self.tags: str = kwargs.get('tags', '')
 
         self.model = kwargs.get('model', None)
         assert self.model is not None
 
     def restore(self):
-        pretrained_model = self.model.load_from_metrics(
-            weights_path=self.weights,
-            tags_csv=self.tags
-        )
+        pretrained_model = self.model.load_from_checkpoint(self.weights)
         pretrained_model.eval()
         pretrained_model.freeze()
         return pretrained_model