From a8882a6896227b83f22f8d5073fb5e1f49f36770 Mon Sep 17 00:00:00 2001 From: alexeykozhevin Date: Thu, 23 Jan 2025 09:19:39 +0000 Subject: [PATCH 1/5] Add ignore_attributes --- batchflow/models/torch/base.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/batchflow/models/torch/base.py b/batchflow/models/torch/base.py index 57c39bcdf..e94ffcf3b 100755 --- a/batchflow/models/torch/base.py +++ b/batchflow/models/torch/base.py @@ -1668,7 +1668,7 @@ def convert_outputs(self, outputs): # Store model def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_openvino=None, - batch_size=None, opset_version=13, pickle_module=dill, **kwargs): + 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). If `use_onnx` is set to True, then the model is converted to ONNX format and stored in a separate file. @@ -1699,6 +1699,8 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op Version of export standard to use. pickle_module : module Module to use for pickling. + ignore_attributes : str ot list, optional + List of attributes to ignore when pickling (e.g. 'optimizer') kwargs : dict Other keyword arguments, passed directly to :func:`torch.save`. """ @@ -1710,6 +1712,11 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op if isinstance(self.model, torch.nn.parallel.DistributedDataParallel): self.model = self.model.module + if isinstance(ignore_attributes, str): + ignore_attributes = [ignore_attributes] + elif ignore_attributes is None: + ignore_attributes = [] + if use_onnx: if batch_size is None: raise ValueError('Specify valid `batch_size`, used for model inference!') @@ -1719,7 +1726,7 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op torch.onnx.export(self.model.eval(), inputs, path_onnx, opset_version=opset_version) # Save the rest of parameters - preserved = set(self.PRESERVE) - set(['model', 'loss', 'optimizer', 'scaler', 'decay']) + preserved = set(self.PRESERVE) - set(ignore_attributes) - set(['model', 'loss', 'optimizer', 'scaler', 'decay']) 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) @@ -1741,12 +1748,13 @@ 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 = set(self.PRESERVE) - set(['model', 'loss', 'optimizer', 'scaler', 'decay']) + preserved = set(self.PRESERVE) - set(ignore_attributes) - set(['model', 'loss', 'optimizer', 'scaler', 'decay']) 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) else: + preserved = set(self.PRESERVE) - set(ignore_attributes) torch.save({item: getattr(self, item) for item in self.PRESERVE}, path, pickle_module=pickle_module, **kwargs) From 562178c0acda2643dd82b5e24237c00be5ea2716 Mon Sep 17 00:00:00 2001 From: arturshagito0 Date: Tue, 28 Jan 2025 11:17:26 +0300 Subject: [PATCH 2/5] Use proper serialization argument (#773) --- batchflow/models/torch/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/batchflow/models/torch/base.py b/batchflow/models/torch/base.py index e94ffcf3b..9611e010d 100755 --- a/batchflow/models/torch/base.py +++ b/batchflow/models/torch/base.py @@ -1755,7 +1755,7 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op else: preserved = set(self.PRESERVE) - set(ignore_attributes) - torch.save({item: getattr(self, item) for item in self.PRESERVE}, + 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): From f178fcf7d95304405fd21279bfb37747b6889c70 Mon Sep 17 00:00:00 2001 From: alexeykozhevin Date: Tue, 28 Jan 2025 13:02:12 +0000 Subject: [PATCH 3/5] Fix linter --- batchflow/models/torch/base.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/batchflow/models/torch/base.py b/batchflow/models/torch/base.py index 9611e010d..5a05684d7 100755 --- a/batchflow/models/torch/base.py +++ b/batchflow/models/torch/base.py @@ -1726,7 +1726,9 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op torch.onnx.export(self.model.eval(), inputs, path_onnx, opset_version=opset_version) # Save the rest of parameters - preserved = set(self.PRESERVE) - set(ignore_attributes) - set(['model', 'loss', 'optimizer', 'scaler', 'decay']) + ignore_attributes = set(ignore_attributes) + set(['model', 'loss', 'optimizer', 'scaler', 'decay']) + preserved = set(self.PRESERVE) - 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) @@ -1748,7 +1750,8 @@ 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 = set(self.PRESERVE) - set(ignore_attributes) - set(['model', 'loss', 'optimizer', 'scaler', 'decay']) + ignore_attributes = set(ignore_attributes) + set(['model', 'loss', 'optimizer', 'scaler', 'decay']) + preserved = set(self.PRESERVE) - 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) From 28d1add7b45787ffa52c10e24e32f2e984e9b511 Mon Sep 17 00:00:00 2001 From: alexeykozhevin Date: Wed, 29 Jan 2025 13:54:39 +0000 Subject: [PATCH 4/5] Change PRESERVE attr --- batchflow/models/torch/base.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/batchflow/models/torch/base.py b/batchflow/models/torch/base.py index 5a05684d7..86b95b490 100755 --- a/batchflow/models/torch/base.py +++ b/batchflow/models/torch/base.py @@ -389,7 +389,7 @@ def callable_init(module): # example of a callable for init 'microbatch_size': 16, # size of microbatches at training } """ - PRESERVE = [ + PRESERVE = set([ 'full_config', 'config', 'model', 'inputs_shapes', 'targets_shapes', 'classes', 'loss', 'optimizer', 'scaler', 'decay', 'decay_step', @@ -397,7 +397,10 @@ def callable_init(module): # example of a callable for init 'iteration', 'last_train_info', 'last_predict_info', 'lr_list', 'syncs', 'decay_iters', '_loss_list', 'loss_list', 'operations' - ] + ]) + + PRESERVE_ONNX = PRESERVE - set(['model', 'loss', 'optimizer', 'scaler', 'decay']) + PRESERVE_OPENVINO = PRESERVE - set(['model', 'loss', 'optimizer', 'scaler', 'decay']) def __init__(self, config=None): if not isinstance(config, (dict, Config)): @@ -1716,6 +1719,7 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op ignore_attributes = [ignore_attributes] elif ignore_attributes is None: ignore_attributes = [] + ignore_attributes = set(ignore_attributes) if use_onnx: if batch_size is None: @@ -1726,8 +1730,7 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op torch.onnx.export(self.model.eval(), inputs, path_onnx, opset_version=opset_version) # Save the rest of parameters - ignore_attributes = set(ignore_attributes) + set(['model', 'loss', 'optimizer', 'scaler', 'decay']) - preserved = set(self.PRESERVE) - ignore_attributes + 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}, @@ -1750,8 +1753,7 @@ 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 - ignore_attributes = set(ignore_attributes) + set(['model', 'loss', 'optimizer', 'scaler', 'decay']) - preserved = set(self.PRESERVE) - ignore_attributes + 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) From 6f7082bef349b657095a94767f1b9e6748744107 Mon Sep 17 00:00:00 2001 From: alexeykozhevin Date: Wed, 29 Jan 2025 13:56:19 +0000 Subject: [PATCH 5/5] Update docstring --- batchflow/models/torch/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/batchflow/models/torch/base.py b/batchflow/models/torch/base.py index 86b95b490..0e7920e69 100755 --- a/batchflow/models/torch/base.py +++ b/batchflow/models/torch/base.py @@ -1702,7 +1702,7 @@ def save(self, path, use_onnx=False, path_onnx=None, use_openvino=False, path_op Version of export standard to use. pickle_module : module Module to use for pickling. - ignore_attributes : str ot list, optional + ignore_attributes : str or iterable, optional List of attributes to ignore when pickling (e.g. 'optimizer') kwargs : dict Other keyword arguments, passed directly to :func:`torch.save`.