Remove all use_xxx and add format=name

Author: igor

batchflow/models/torch/base.py

@@ -1669,8 +1669,7 @@ def convert_outputs(self, outputs):
 
 
     # Store model
-    def save(self, path=None, use_onnx=False, path_onnx=None, use_openvino=False, path_openvino=None,
-             use_safetensors=False, path_safetensors=None, pickle_metadata=False,
+    def save(self, path, format="pt", pickle_metadata=False,
              batch_size=None, opset_version=13, pickle_module=dill, ignore_attributes=None, **kwargs):
         """ Save underlying PyTorch model along with meta parameters (config, device spec, etc).
 
@@ -1681,7 +1680,7 @@ def save(self, path=None, use_onnx=False, path_onnx=None, use_openvino=False, pa
 
         Parameters
         ----------
-        path : Optional[str]
+        path : str
             Path to a file where the model data will be stored.
         use_onnx: bool
             Whether to store model in ONNX format.
@@ -1709,14 +1708,11 @@ def save(self, path=None, use_onnx=False, path_onnx=None, use_openvino=False, pa
         kwargs : dict
             Other keyword arguments, passed directly to :func:`torch.save`.
         """
-
-        if path is None:
-            dirname = os.getcwd()
-            path = os.path.join(dirname, "model.pt")
-        else:
-            dirname = os.path.dirname(path)
-            if dirname and not os.path.exists(dirname):
-                os.makedirs(dirname)
+        available_formats = ("pt", "onnx", "openvino", "safetensors")
+        assert format in available_formats, f"Format must be in {available_formats}"
+        dirname = os.path.dirname(path)
+        if dirname and not os.path.exists(dirname):
+            os.makedirs(dirname)
 
         # Unwrap DDP model
         if isinstance(self.model, torch.nn.parallel.DistributedDataParallel):
@@ -1728,12 +1724,12 @@ def save(self, path=None, use_onnx=False, path_onnx=None, use_openvino=False, pa
             ignore_attributes = []
         ignore_attributes = set(ignore_attributes)
 
-        if use_onnx:
+        if format == "onnx":
             if batch_size is None:
                 raise ValueError('Specify valid `batch_size`, used for model inference!')
 
             inputs = self.make_placeholder_data(batch_size=batch_size, unwrap=False)
-            path_onnx = path_onnx or os.path.join(dirname, "model.onnx")
+            path_onnx = path if path.endswith(".onnx") else os.path.join(dirname, "model.onnx")
             torch.onnx.export(self.model.eval(), inputs, path_onnx, opset_version=opset_version)
 
             if pickle_metadata:
@@ -1744,10 +1740,10 @@ def save(self, path=None, use_onnx=False, path_onnx=None, use_openvino=False, pa
                 torch.save({'onnx': True, 'path_onnx': path_onnx, 'onnx_batch_size': batch_size, **preserved_dict},
                         path, pickle_module=pickle_module, **kwargs)
 
-        elif use_openvino:
+        elif format == "openvino":
             import openvino as ov
 
-            path_openvino = path_openvino or os.path.join(dirname, "model.openvino")
+            path_openvino = path if path.endswith(".openvino") else os.path.join(dirname, "model.openvino")
             if os.path.splitext(path_openvino)[-1] == '':
                 path_openvino = f'{path_openvino}.xml'
 
@@ -1767,11 +1763,11 @@ def save(self, path=None, use_onnx=False, path_onnx=None, use_openvino=False, pa
                 torch.save({'openvino': True, 'path_openvino': path_openvino, **preserved_dict},
                         path, pickle_module=pickle_module, **kwargs)
 
-        elif use_safetensors:
+        elif format == "safetensors":
             from safetensors.torch import save_file
             state_dict = self.model.state_dict()
 
-            path_safetensors = path_safetensors or os.path.join(dirname, "model.safetensors")
+            path_safetensors = path if path.endswith(".safetensors") else os.path.join(dirname, "model.safetensors")
             save_file(state_dict, path_safetensors)
 
             preserved = self.PRESERVE_SAFETENSORS - ignore_attributes