Added zlib compression helper file and cleaned up zfp helper

Author: alex-grindrod

scripts/zfp_compress.py

@@ -6,9 +6,24 @@
 import os
 import zfpy
 import copy
-import zipfile
-import numpy as np
-import io
+
+
+def zfp_compress(model, name="generic_model_type", tolerance=1e-3):
+    zfp_model = copy.deepcopy(model)
+
+    org_size = save_pkl(f'models/{name}.pkl', model.state_dict())
+    print(f"Original Size: {org_size}")
+
+    compressed_params = apply_zfp(model, tolerance)
+    compressed_size = save_pkl(f"models/{name}.pkl", compressed_params)
+    print(f"Compressed Size: {compressed_size}")
+
+    load_and_decompress(f"models/{name}.pkl", zfp_model)
+
+    return zfp_model, org_size, compressed_size
+
+
+#HELPERS
 
 def save_pkl(file: str, state_dict: dict) -> float:
     with open(file, 'wb') as f:
@@ -35,43 +50,4 @@ def load_and_decompress(file: str, model) -> None:
             decompressed = zfpy.decompress_numpy(params[name])
             param.data = torch.tensor(decompressed).to(device)
 
-def zfp_compress(model, name, tolerance=1e-3):
-    zfp_model = copy.deepcopy(model)
-
-    org_size = save_pkl(f'models/{name}.pkl', model.state_dict())
-    print(f"Original Size: {org_size}")
-
-    compressed_params = apply_zfp(model, tolerance)
-    compressed_size = save_pkl(f"models/{name}.pkl", compressed_params)
-    print(f"Compressed Size: {compressed_size}")
-
-    load_and_decompress(f"models/{name}.pkl", zfp_model)
-
-    return zfp_model, org_size, compressed_size
-
-
-
-
-# def apply_zip(model):
-#     params = dict()
-#     for name, param in model.named_parameters():
-#         if param.requires_grad:
-#             param_cpu = param.cpu().detach().numpy()
-
-#             array_buffer = io.BytesIO()
-#             np.save(array_buffer, param_cpu)
-
-
-#             # params[name] = zfpy.compress_numpy(param_cpu, tolerance=tolerance)
-#     return params
-
-# def zip_compress(model, name):
-#     zip_model = copy.deepcopy(model)
-
-#     pkl_path = f'models/{name}.pkl'
-#     org_size = save_pkl(pkl_path, model.state_dict())
-
-
-
 
-    

scripts/zlib_compress.py

@@ -0,0 +1,40 @@
+import copy
+from scripts.zfp_compress import save_pkl, load_pkl
+import torch
+import zlib
+import numpy as np
+
+
+def zlib_compress(model, name="generic_model_type"):
+    zlib_model = copy.deepcopy(model)
+    org_size = save_pkl(f'models/{name}.pkl', model.state_dict())
+    print(f"Original Size: {org_size}")
+
+    compressed_params, shapes = apply_zlib(model)
+    compressed_size = save_pkl(f"models/{name}.pkl", compressed_params)
+    print(f"Compressed Size: {compressed_size}")
+
+    load_and_decompress(f"models/{name}.pkl", zlib_model, shapes)
+
+    return zlib_model, org_size, compressed_size
+
+def apply_zlib(model):    
+    params = {}
+    shapes = {}
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            param_cpu = param.cpu().detach().numpy()
+            params[name] = zlib.compress(param_cpu)  
+            shapes[name] = param.size()
+    return params, shapes
+
+
+def load_and_decompress(file: str, model, shapes) -> None:
+    device = torch.device('cpu')
+    params = load_pkl(file)
+    for name, param in model.named_parameters():
+        if name in params:
+            decompressed_bytes = zlib.decompress(params[name])
+            decompressed = np.frombuffer(decompressed_bytes, dtype=np.float32)
+            shape = shapes[name]
+            param.data = torch.tensor(decompressed.reshape(shape)).to(device)