added weights_only

myresources/crocodile/deeplearning.py

@@ -489,7 +489,8 @@ def save_class(self, weights_only: bool = True, version: str = 'v0', strict: boo
  try: generate_readme(get_hp_save_dir(hp), obj=self.__class__, desc=desc)
  except Exception as ex: print(ex) # often fails because model is defined in main during experiments.
  save_dir = get_hp_save_dir(hp).joinpath(f'{"weights" if weights_only else "model"}_save_{version}')
- if weights_only: self.save_weights(save_dir.create())
+ if weights_only:
+ self.save_weights(save_dir.create())
  else:
  self.save_model(save_dir)
 

myresources/crocodile/deeplearning_pytorch_template.py

@@ -96,7 +96,7 @@ def main():
  m.save_weights(save_dir=save_dir)
  # save_onnx(save_dir=save_dir, model=m.model, dummy_ip=x)
 
- m1 = BaseModel.load_model(save_dir=save_dir, map_location=None)
+ m1 = BaseModel.load_model(save_dir=save_dir, map_location=None, weights_only=True)
  m_init = My2LayerNN(hp=hp)
  # m_base = BaseModel(m_init, optimizer=optimizer, loss=loss, metrics=[])
  m2 = BaseModel.load_weights(model=m_init, save_dir=save_dir, map_location=None)

myresources/crocodile/deeplearning_torch.py

@@ -92,10 +92,10 @@ def summary(model: nn.Module, detailed: bool = False):
 
  def save_model(self, save_dir: P): t.save(self.model, save_dir.joinpath("model.pth"))
  @staticmethod
- def load_model(save_dir: P, map_location: Union[str, Device, None]):
+ def load_model(save_dir: P, map_location: Union[str, Device, None], weights_only: bool):
  if map_location is None and t.cuda.is_available():
  map_location = "cpu"
- model: nn.Module = t.load(save_dir.joinpath("model.pth"), map_location=map_location, weights_only=True) # type: ignore
+ model: nn.Module = t.load(save_dir.joinpath("model.pth"), map_location=map_location, weights_only=weights_only) # type: ignore
  model.eval()
  import traceback
  try: