Add save in safetensors format

batchflow/models/torch/base.py

@@ -15,7 +15,6 @@
 from torch import nn
 from torch.optim.swa_utils import AveragedModel, SWALR
 
-
 from sklearn.decomposition import PCA
 
 from ...utils_import import make_delayed_import
@@ -400,6 +399,7 @@ def callable_init(module): # example of a callable for init
 
     PRESERVE_ONNX = PRESERVE - set(['model', 'loss', 'optimizer', 'scaler', 'decay'])
     PRESERVE_OPENVINO = PRESERVE - set(['model', 'loss', 'optimizer', 'scaler', 'decay'])
+    PRESERVE_SAFETENSORS = PRESERVE - set(['model', 'loss', 'optimizer', 'scaler', 'decay'])
 
     def __init__(self, config=None):
         if not isinstance(config, (dict, Config)):
@@ -987,12 +987,16 @@ def transfer_from_device(self, data, force_float32_dtype=True):
 
     def model_to_device(self, model=None):
         """ Put model on device(s). If needed, apply DataParallel wrapper. """
-        model = model if model is not None else self.model
+        model_ = model if model is not None else self.model
 
         if len(self.devices) > 1:
-            model = nn.DataParallel(model, self.devices)
+            model_ = nn.DataParallel(model_, self.devices)
         else:
-            model = model.to(self.device)
+            model_ = model_.to(self.device)
+
+        if model is None:
+            self.model = model_
+        return model_
 
 
     # Apply model to train/predict on given data
@@ -1669,7 +1673,7 @@ def convert_outputs(self, outputs):
 
 
     # Store model
-    def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_openvino=None,
+    def save(self, path, fmt=None, pickle_metadata=True,
              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).
 
@@ -1682,17 +1686,10 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op
         ----------
         path : str
             Path to a file where the model data will be stored.
-        use_onnx: bool
-            Whether to store model in ONNX format.
-        path_onnx : str, optional
-            Used only if `use_onnx` is True.
-            If provided, then path to store the ONNX model; default `path_onnx` is `path` with '_onnx' postfix.
-        use_openvino: bool
-            Whether to store model as openvino xml file.
-        path_openvino : str, optional
-            Used only if `use_openvino` is True.
-            If provided, then path to store the openvino model; default `path_openvino` is `path` with '_openvino'
-            postfix.
+        fmt: Optional[str]
+            Weights format. Available formats: "pt", "onnx", "openvino", "safetensors"
+        pickle_metadata: bool
+            Whether make pickle with metadata
         batch_size : int, optional
             Used only if `use_onnx` is True.
             Fixed batch size of the ONNX module. This is the only viable batch size for this model after loading.
@@ -1706,6 +1703,14 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op
         kwargs : dict
             Other keyword arguments, passed directly to :func:`torch.save`.
         """
+        available_formats = ("pt", "onnx", "openvino", "safetensors")
+
+        if fmt is None:
+            fmt = os.path.splitext(path)[-1][1:]
+
+        if fmt not in available_formats:
+            raise ValueError(f"fmt must be in {available_formats}")
+
         dirname = os.path.dirname(path)
         if dirname and not os.path.exists(dirname):
             os.makedirs(dirname)
@@ -1720,25 +1725,27 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op
             ignore_attributes = []
         ignore_attributes = set(ignore_attributes)
 
-        if use_onnx:
+        if fmt == "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 (path + '_onnx')
+
+            path_onnx = path if not pickle_metadata else os.path.splitext(path)[0] + ".onnx"
             torch.onnx.export(self.model.eval(), inputs, path_onnx, opset_version=opset_version)
 
-            # Save the rest of parameters
-            preserved = self.PRESERVE_ONNX - ignore_attributes
+            if pickle_metadata:
+                # Save the rest of parameters
+                preserved = self.PRESERVE_ONNX - ignore_attributes
 
-            preserved_dict = {item: getattr(self, item) for item in preserved}
-            torch.save({'onnx': True, 'path_onnx': path_onnx, 'onnx_batch_size': batch_size, **preserved_dict},
-                       path, pickle_module=pickle_module, **kwargs)
+                preserved_dict = {item: getattr(self, item) for item in preserved}
+                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 fmt == "openvino":
             import openvino as ov
 
-            path_openvino = path_openvino or (path + '_openvino')
+            path_openvino = path if not pickle_metadata else os.path.splitext(path)[0] + ".openvino"
             if os.path.splitext(path_openvino)[-1] == '':
                 path_openvino = f'{path_openvino}.xml'
 
@@ -1751,18 +1758,33 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op
 
             ov.save_model(model, output_model=path_openvino)
 
-            # Save the rest of parameters
-            preserved = self.PRESERVE_OPENVINO - ignore_attributes
+            if pickle_metadata:
+                # Save the rest of parameters
+                preserved = self.PRESERVE_OPENVINO - ignore_attributes
+                preserved_dict = {item: getattr(self, item) for item in preserved}
+                torch.save({'openvino': True, 'path_openvino': path_openvino, **preserved_dict},
+                        path, pickle_module=pickle_module, **kwargs)
+
+        elif fmt == "safetensors":
+            from safetensors.torch import save_file
+            state_dict = self.model.state_dict()
+
+            path_safetensors = path if not pickle_metadata else os.path.splitext(path)[0] + ".safetensors"
+            save_file(state_dict, path_safetensors)
+
+            preserved = self.PRESERVE_SAFETENSORS - ignore_attributes
             preserved_dict = {item: getattr(self, item) for item in preserved}
-            torch.save({'openvino': True, 'path_openvino': path_openvino, **preserved_dict},
-                       path, pickle_module=pickle_module, **kwargs)
+
+            if pickle_metadata:
+                torch.save({'safetensors': True, 'path_safetensors': path_safetensors, **preserved_dict},
+                        path, pickle_module=pickle_module, **kwargs)
 
         else:
             preserved = set(self.PRESERVE) - set(ignore_attributes)
             torch.save({item: getattr(self, item) for item in preserved},
                        path, pickle_module=pickle_module, **kwargs)
 
-    def load(self, file, make_infrastructure=False, mode='eval', pickle_module=dill, **kwargs):
+    def load(self, file, fmt=None, make_infrastructure=False, mode='eval', pickle_module=dill, **kwargs):
         """ Load a torch model from a file.
 
         If the model was saved in ONNX format (refer to :meth:`.save` for more info), we fix the microbatch size
@@ -1772,6 +1794,8 @@ def load(self, file, make_infrastructure=False, mode='eval', pickle_module=dill,
         ----------
         file : str, PathLike, io.Bytes
             a file where a model is stored.
+        fmt: optional str
+            Weights format. Available formats: "pt", "onnx", "openvino", "safetensors"
         make_infrastructure : bool
             Whether to re-create model loss, optimizer, scaler and decay.
         mode : str
@@ -1793,6 +1817,53 @@ def load(self, file, make_infrastructure=False, mode='eval', pickle_module=dill,
         else:
             self._parse_devices()
 
+        if isinstance(file, str):
+            if fmt == "safetensors" or (fmt is None and file.endswith(".safetensors")):
+                from safetensors.torch import load_file
+                state_dict = load_file(file)
+
+                inputs = self.make_placeholder_data(to_device=True)
+                with torch.no_grad():
+                    self.model = Network(inputs=inputs, config=self.config, device=self.device)
+
+                self.model.load_state_dict(state_dict)
+
+                self.model_to_device()
+
+                if make_infrastructure:
+                    self.make_infrastructure()
+
+                self.set_model_mode(mode)
+
+                return
+
+            if fmt == "onnx" or (fmt is None and file.endswith(".onnx")):
+                try:
+                    from onnx2torch import convert
+                except ImportError as e:
+                    raise ImportError('Loading model, stored in ONNX format, requires `onnx2torch` library.') from e
+
+                model = convert(file).eval()
+                self.model = model
+
+                self.model_to_device()
+
+                if make_infrastructure:
+                    self.make_infrastructure()
+
+                self.set_model_mode(mode)
+
+                return
+
+            if fmt == "openvino" or (fmt is None and file.endswith(".openvino")):
+                model = OVModel(model_path=file, **model_load_kwargs)
+                self.model = model
+
+                self._loaded_from_openvino = True
+                self.disable_training = True
+
+                return
+
         kwargs['map_location'] = self.device
 
         # Load items from disk storage and set them as insance attributes
@@ -1828,6 +1899,11 @@ def load(self, file, make_infrastructure=False, mode='eval', pickle_module=dill,
                 self._loaded_from_onnx = True
                 self.disable_training = True
 
+            if "safetensors" in checkpoint:
+                from safetensors.torch import load_file
+                state_dict = load_file(checkpoint['path_safetensors'], device=device)
+                self.model.load_state_dict(state_dict)
+
             self.model_to_device()
 
             if make_infrastructure:

batchflow/tests/research_test.py

@@ -570,6 +570,7 @@ def f(a):
         assert research.results.df.iloc[0].a == f(2)
         assert research.results.df.iloc[0].b == f(3)
 
+    @pytest.mark.slow
     @pytest.mark.parametrize('dump_results', [False, True])
     @pytest.mark.parametrize('redirect_stdout', [True, 0, 1, 2, 3])
     @pytest.mark.parametrize('redirect_stderr', [True, 0, 1, 2, 3])

pyproject.toml

@@ -74,6 +74,10 @@ telegram = [
     "pillow>=9.4,<11.0",
 ]
 
+safetensors = [
+    "safetensors>=0.5.3",
+]
+
 other = [
     "urllib3>=1.25"
 ]