Merge pull request mala-project#660 from karanprime/develop

Added json functionality for data scaling

Author: RandomDefaultUser

mala/datahandling/data_scaler.py

@@ -4,6 +4,8 @@
 import numpy as np
 import torch
 import torch.distributed as dist
+import json
+import io
 
 from mala.common.parameters import printout
 from mala.common.parallelizer import parallel_warn
@@ -517,7 +519,7 @@ def inverse_transform(self, scaled, copy=False, as_numpy=False):
         else:
             return unscaled
 
-    def save(self, filename, save_format="pickle"):
+    def save(self, filename, save_format="json"):
         """
         Save the Scaler object so that it can be accessed again later.
 
@@ -527,23 +529,58 @@ def save(self, filename, save_format="pickle"):
             File in which the parameters will be saved.
 
         save_format :
-            File format which will be used for saving.
+            File format which will be used for saving. Default is "json".  
+            Pickle format is deprecated and will be removed in future versions.
         """
         # If we use ddp, only save the network on root.
         if self.use_ddp:
             if dist.get_rank() != 0:
                 return
-        if save_format == "pickle":
+        
+        filename_format = filename.rsplit(".", 1)[1]
+        if save_format == "pickle" or filename_format == "pkl": # similar to "normal" string warning
+            parallel_warn(
+            "Pickle format is deprecated and will be removed in future versions. "
+            "Please use JSON format instead.",
+            min_verbosity=0,
+            category=FutureWarning,
+        )
             with open(filename, "wb") as handle:
                 pickle.dump(self, handle, protocol=4)
+        elif save_format == "json" or filename_format == "json":
+            # saving tensors as lists for json
+            # if scale_normal is used, it will be converted to scale_minmax
+            data_dict = {
+                "typestring": self.typestring,
+                "use_ddp": self.use_ddp,
+                "scale_standard": self.scale_standard,
+                "scale_minmax": (self.scale_minmax 
+                               if hasattr(self, "scale_minmax") 
+                               else self.scale_normal),
+                "feature_wise": self.feature_wise,
+                "cantransform": self.cantransform,
+                "means": self.means.tolist() if hasattr(self.means, "tolist") else [],
+                "stds": self.stds.tolist() if hasattr(self.stds, "tolist") else [],
+                "maxs": self.maxs.tolist() if hasattr(self.maxs, "tolist") else [],
+                "mins": self.mins.tolist() if hasattr(self.mins, "tolist") else [],
+                "total_mean": float(self.total_mean),
+                "total_std": float(self.total_std),
+                "total_max": float(self.total_max),
+                "total_min": float(self.total_min),
+                "total_data_count": self.total_data_count
+            }
+            
+            with open(filename, "w") as handle:
+                json.dump(data_dict, handle, indent=4)
+                
         else:
             raise Exception("Unsupported parameter save format.")
 
     @classmethod
-    def load_from_file(cls, file, save_format="pickle"):
+    def load_from_file(cls, file, save_format="json", auto_convert=True):
         """
         Load a saved Scaler object.
-
+        
         Parameters
         ----------
         file : string or ZipExtFile
@@ -552,17 +589,73 @@ def load_from_file(cls, file, save_format="pickle"):
         save_format :
             File format which was used for saving.
 
+        auto_convert : bool
+            If True and loading from pickle format, automatically save as JSON for future use.
+
         Returns
         -------
         data_scaler : DataScaler
             DataScaler which was read from the file.
         """
-        if save_format == "pickle":
+        if isinstance(file, str):
+            filename = file
+        elif hasattr(file, 'name'): # getting fname from zip file
+            filename = file.name
+        else:
+            raise Exception("File must be either a string path or a ZipFile object")
+        
+        filename_format = filename.rsplit(".", 1)[1]
+
+        if save_format == "pickle" or filename_format == "pkl":
+            parallel_warn(
+                "Loading from pickle format is deprecated and will be removed in future versions. "
+                "Please convert your files to JSON format.",
+                min_verbosity=0,
+                category=FutureWarning,
+            )
             if isinstance(file, str):
                 loaded_scaler = pickle.load(open(file, "rb"))
-            else:
+
+                if auto_convert:
+                    json_file_path = filename.rsplit(".", 1)[0] + ".json"
+                    loaded_scaler.save(json_file_path, save_format="json")
+                    
+
+            elif hasattr(file, 'name'):
                 loaded_scaler = pickle.load(file)
-        else:
-            raise Exception("Unsupported parameter save format.")
 
+            parallel_warn(
+                        "Pickle file has been automatically converted to JSON format.",
+                        min_verbosity=0,
+                        category=FutureWarning,
+                    )            
+            
+        elif save_format == "json" or filename_format == "json":
+            if isinstance(file, str):
+                with open(file, "r") as handle:
+                    data_dict = json.load(handle)
+            elif hasattr(file, 'name'):
+                text_handle = io.TextIOWrapper(file, encoding="utf-8")
+                data_dict = json.load(text_handle)
+            
+                loaded_scaler = cls(data_dict["typestring"], data_dict["use_ddp"])
+                
+                loaded_scaler.scale_standard = data_dict["scale_standard"]
+                loaded_scaler.scale_minmax = data_dict["scale_minmax"]
+                loaded_scaler.feature_wise = data_dict["feature_wise"]
+                loaded_scaler.cantransform = data_dict["cantransform"]
+                
+                loaded_scaler.means = torch.tensor(data_dict["means"])
+                loaded_scaler.stds = torch.tensor(data_dict["stds"])
+                loaded_scaler.maxs = torch.tensor(data_dict["maxs"])
+                loaded_scaler.mins = torch.tensor(data_dict["mins"])
+                
+                loaded_scaler.total_mean = torch.tensor(data_dict["total_mean"])
+                loaded_scaler.total_std = torch.tensor(data_dict["total_std"])
+                loaded_scaler.total_max = torch.tensor(data_dict["total_max"])
+                loaded_scaler.total_min = torch.tensor(data_dict["total_min"])
+                loaded_scaler.total_data_count = data_dict["total_data_count"]
+        else:
+            raise Exception("Unsupported parameter save format. Use 'json' or 'pickle'.")
+            
         return loaded_scaler

mala/network/runner.py

@@ -1,9 +1,11 @@
 """Runner class for running networks."""
 
 import os
+import tempfile
 from zipfile import ZipFile, ZIP_STORED
 
 from mala.common.parallelizer import printout
+from mala.common.parallelizer import parallel_warn
 
 import numpy as np
 import torch
@@ -580,8 +582,8 @@ def save_run(
         # performed on rank 0.
         if get_rank() == 0:
             model_file = run_name + ".network.pth"
-            iscaler_file = run_name + ".iscaler.pkl"
-            oscaler_file = run_name + ".oscaler.pkl"
+            iscaler_file = run_name + ".iscaler.json"
+            oscaler_file = run_name + ".oscaler.json"
             params_file = run_name + ".params.json"
             if save_runner:
                 optimizer_file = run_name + ".optimizer.pth"
@@ -632,6 +634,7 @@ def load_run(
         path="./",
         zip_run=True,
         params_format="json",
+        scalers_format="json",
         load_runner=True,
         prepare_data=False,
         load_with_mpi=None,
@@ -669,6 +672,10 @@ def load_run(
             Can be "json" or "pkl", depending on what was saved by the model.
             Default is "json".
 
+        scalers_format: str
+            Can be "json" or "pkl", depending on what was saved by the model.
+            Default is "json".
+
         load_runner : bool
             If True, a Runner object will be created/loaded for further use.
 
@@ -719,26 +726,44 @@ def load_run(
         loaded_info = None
         if zip_run is True:
             loaded_network = run_name + ".network.pth"
-            loaded_iscaler = run_name + ".iscaler.pkl"
-            loaded_oscaler = run_name + ".oscaler.pkl"
+            loaded_iscaler = run_name + ".iscaler." + scalers_format
+            loaded_oscaler = run_name + ".oscaler." + scalers_format
             loaded_params = run_name + ".params." + params_format
             loaded_info = run_name + ".info.json"
 
+            iscale_pickle_flag = False
+            oscale_pickle_flag = False
+
             zip_path = os.path.join(path, run_name + ".zip")
             with ZipFile(zip_path, "r") as zip_obj:
                 loaded_params = zip_obj.open(loaded_params)
                 loaded_network = zip_obj.open(loaded_network)
-                loaded_iscaler = zip_obj.open(loaded_iscaler)
-                loaded_oscaler = zip_obj.open(loaded_oscaler)
+                
+                # If json scaler files not found, try pickle format
+                try:
+                    loaded_iscaler = zip_obj.open(loaded_iscaler)
+                except KeyError:
+                    iscale_pickle_flag = True
+                    loaded_iscaler = zip_obj.open(loaded_iscaler.replace(".json", ".pkl"))   
+                try:
+                    loaded_oscaler = zip_obj.open(loaded_oscaler)
+                except KeyError:
+                    oscale_pickle_flag = True
+                    loaded_oscaler = zip_obj.open(loaded_oscaler.replace(".json", ".pkl"))
+
                 if loaded_info in zip_obj.namelist():
                     loaded_info = zip_obj.open(loaded_info)
                 else:
                     loaded_info = None
 
         else:
             loaded_network = os.path.join(path, run_name + ".network.pth")
-            loaded_iscaler = os.path.join(path, run_name + ".iscaler.pkl")
-            loaded_oscaler = os.path.join(path, run_name + ".oscaler.pkl")
+            loaded_iscaler = os.path.join(
+                path, run_name + ".iscaler." + scalers_format
+            )
+            loaded_oscaler = os.path.join(
+                path, run_name + ".oscaler." + scalers_format
+            )
             loaded_params = os.path.join(
                 path, run_name + ".params." + params_format
             )
@@ -772,6 +797,37 @@ def load_run(
         loaded_network = Network.load_from_file(loaded_params, loaded_network)
         loaded_iscaler = DataScaler.load_from_file(loaded_iscaler)
         loaded_oscaler = DataScaler.load_from_file(loaded_oscaler)
+
+        # only on rank 0, if pickle scaler files are found,
+        # add their json versions to the existing zip file
+        if get_rank() == 0 and (zip_run and (iscale_pickle_flag or oscale_pickle_flag)):
+            parallel_warn(
+                        "Pickle file has been automatically converted to JSON format.",
+                        min_verbosity=0,
+                        category=FutureWarning,
+                    )
+            with tempfile.TemporaryDirectory() as temp_dir:
+                with ZipFile(zip_path, 'r') as zip_read:
+                    zip_read.extractall(temp_dir)
+
+                iscaler_file = run_name + ".iscaler.json"
+                oscaler_file = run_name + ".oscaler.json"
+                
+                loaded_iscaler.save(os.path.join(temp_dir, iscaler_file))
+                loaded_oscaler.save(os.path.join(temp_dir, oscaler_file))
+
+                temp_zip_path = zip_path + ".temp"
+                with ZipFile(temp_zip_path, 'w') as zip_write:
+                    for foldername, subfolders, filenames in os.walk(temp_dir):
+                        for filename in filenames:
+                            file_path = os.path.join(foldername, filename)
+                            arcname = os.path.relpath(file_path, temp_dir)
+                            zip_write.write(file_path, arcname)
+
+                os.replace(temp_zip_path, zip_path)
+            
+
+
         new_datahandler = DataHandler(
             loaded_params,
             input_data_scaler=loaded_iscaler,

test/README.md

@@ -1,7 +1,7 @@
 # Test directory
 
-Different tests that can be rerun at any time to make sure a certain function or idea stillm works.
+Different tests that can be rerun at any time to make sure a certain function or idea still works.
 
 ## tensor_memory.py
 
-Verifies that the way we create torch tensors from numpy arrays is in fact by referencing, and not by copying. 
+Verifies that the way we create torch tensors from numpy arrays is in fact by referencing, and not by copying.