RRAEsTorch 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

RRAEsTorch/training_classes/training_classes.py

@@ -13,7 +13,6 @@ import os
 import time
 import dill
 import shutil
-from RRAEsTorch.wrappers import vmap_wrap, norm_wrap
 from functools import partial
 from RRAEsTorch.trackers import (
     Null_Tracker,
@@ -26,6 +25,8 @@ from prettytable import PrettyTable
 import torch
 from torch.utils.data import TensorDataset, DataLoader
 
+from RRAEsTorch.utilities import get_basis
+
 class Circular_list:
     """
     Creates a list of fixed size.
@@ -157,40 +158,17 @@ class Trainor_class:
         model_cls=None,
         folder="",
         file=None,
-        out_train=None,
-        norm_in="None",
-        norm_out="None",
-        methods_map=["__call__"],
-        methods_norm_in=["__call__"],
-        methods_norm_out=["__call__"],
-        call_map_count=1,
-        call_map_axis=-1,
         **kwargs,
     ):
         if model_cls is not None:
             orig_model_cls = model_cls
-            model_cls = vmap_wrap(orig_model_cls, call_map_axis, call_map_count, methods_map)
-            model_cls = norm_wrap(model_cls, in_train, norm_in, None, out_train, norm_out, None, methods_norm_in, methods_norm_out)
             self.model = model_cls(**kwargs)
-            params_in = self.model.params_in
-            params_out = self.model.params_out
         else:
             orig_model_cls = None
-            params_in = None
-            params_out = None
 
         self.all_kwargs = {
             "kwargs": kwargs,
-            "params_in": params_in,
-            "params_out": params_out,
-            "norm_in": norm_in,
-            "norm_out": norm_out,
-            "call_map_axis": call_map_axis,
-            "call_map_count": call_map_count,
-            "orig_model_cls": orig_model_cls,
-            "methods_map": methods_map,
-            "methods_norm_in": methods_norm_in,
-            "methods_norm_out": methods_norm_out,
+            "orig_model_cls": orig_model_cls
         }
 
         self.folder = folder
@@ -229,7 +207,8 @@ class Trainor_class:
         save_losses=False,
         input_val=None,
         output_val=None,
-        latent_size=0
+        latent_size=0,
+        device="cpu"
     ):
         assert isinstance(input, torch.Tensor), "Input should be a torch tensor"
         assert isinstance(output, torch.Tensor), "Output should be a torch tensor"
@@ -301,7 +280,12 @@ class Trainor_class:
             all_losses = []
 
         
-        
+        model = model.to(device)
+
+        if input_val is not None:
+            dataset_val = TensorDataset(input_val, output_val, torch.arange(0, input_val.shape[0], 1))
+            dataloader_val = DataLoader(dataset_val, batch_size=input_val.shape[0], shuffle=False)
+
         # Outler Loop
         for steps, lr, batch_size in zip(step_st, lr_st, batch_size_st):
             try:
@@ -310,15 +294,13 @@ class Trainor_class:
                     filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3
                 )
 
-                if (batch_size > input.shape[-1]) or batch_size == -1:
-                    print(f"Setting batch size to: {input.shape[-1]}")
-                    batch_size = input.shape[-1]
+                if (batch_size > input.shape[0]) or batch_size == -1:
+                    print(f"Setting batch size to: {input.shape[0]}")
+                    batch_size = input.shape[0]
 
                 # Inner loop (batch)
-                inputT = input.permute(*range(input.ndim - 1, -1, -1))
-                outputT = output.permute(*range(output.ndim - 1, -1, -1))
 
-                dataset = TensorDataset(inputT, outputT, torch.arange(0, input.shape[-1], 1))
+                dataset = TensorDataset(input, output, torch.arange(0, input.shape[0], 1))
                 dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
                 data_iter = iter(dataloader)
 
@@ -331,15 +313,16 @@ class Trainor_class:
                         input_b, out_b, idx_b = next(data_iter)
     
                     start_time = time.perf_counter()             # Start time
-                    input_b = input_b.permute(*range(input_b.ndim - 1, -1, -1))
-                    out_b = self.model.norm_out.default(None, pre_func_out(out_b))    # Pre-process batch out values
-                    out_b = out_b.permute(*range(out_b.ndim - 1, -1, -1))
+                    out_b = pre_func_out(out_b)   # Pre-process batch out values
                     input_b = pre_func_inp(input_b)              # Pre-process batch input values 
-                    epsilon = eps_fn(latent_size, input_b.shape[-1])
+                    epsilon = eps_fn(latent_size, input_b.shape[0])
 
                     step_kwargs = merge_dicts(loss_kwargs, track_params)
 
                     # Compute loss
+                    input_b = input_b.to(device)
+                    out_b = out_b.to(device)
+                    
                     loss, model, optimizer_tr, (aux, extra_track) = make_step(
                                                                 model,
                                                                 input_b,
@@ -351,12 +334,14 @@ class Trainor_class:
                                                             )
                     
                     if input_val is not None:
-    
-
-                        idx = np.arange(input_val.shape[-1])
-                        val_loss, _ = loss_fun(
-                            model, input_val, output_val, idx=idx, epsilon=None, **step_kwargs
-                        )
+                        val_loss = []
+                        for input_vb, out_vb, idx_b in dataloader_val:
+                            out_vb = pre_func_out(out_vb) 
+                            val_loss_batch = loss_fun(
+                                model, input_vb.to(device), out_vb.to(device), idx=idx_b, epsilon=None, **step_kwargs
+                            )[0]
+                            val_loss.append(val_loss_batch.item())
+                        val_loss = sum(val_loss) / len(val_loss)
                         aux["val_loss"] = val_loss
                     else:
                         aux["val_loss"] = None
@@ -496,18 +481,16 @@ class Trainor_class:
                 hasattr(self, "batch_size") or batch_size is not None
             ), "You should either provide a batch_size or fit the model first."
 
-            x_train_oT = x_train_o.permute(*range(x_train_o.ndim - 1, -1, -1))
-            dataset = TensorDataset(x_train_oT)
+            dataset = TensorDataset(x_train_o)
             batch_size = self.batch_size if batch_size is None else batch_size
             dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
             
             pred = []
             for x_b in dataloader:
-                x_bT = x_b[0].permute(*range(x_b[0].ndim - 1, -1, -1))
-                pred_batch = call_func(x_bT)
+                pred_batch = call_func(x_b[0])
                 pred.append(pred_batch)
 
-            y_pred_train_o = torch.concatenate(pred, axis=-1)
+            y_pred_train_o = torch.concatenate(pred)
 
             self.error_train_o = (
                 torch.linalg.norm(y_pred_train_o - y_train_o)
@@ -516,8 +499,8 @@ class Trainor_class:
             )
             print("Train error on original output: ", self.error_train_o)
 
-            y_pred_train = self.model.norm_out.default(None, y_pred_train_o)
-            y_train = self.model.norm_out.default(None, y_train_o)
+            y_pred_train = pre_func_out(y_pred_train_o)
+            y_train = pre_func_out(y_train_o)
             self.error_train = (
                 torch.linalg.norm(y_pred_train - y_train) / torch.linalg.norm(y_train) * 100
             )
@@ -525,16 +508,14 @@ class Trainor_class:
 
         if x_test_o is not None:
             y_test_o = pre_func_out(y_test_o)
-            x_test_oT = x_test_o.permute(*range(x_test_o.ndim - 1, -1, -1))
-            dataset = TensorDataset(x_test_oT)
+            dataset = TensorDataset(x_test_o)
             batch_size = self.batch_size if batch_size is None else batch_size
             dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
             pred = []
             for x_b in dataloader:
-                x_bT = x_b[0].permute(*range(x_b[0].ndim - 1, -1, -1))
-                pred_batch = call_func(x_bT)
+                pred_batch = call_func(x_b[0])
                 pred.append(pred_batch)
-            y_pred_test_o = torch.concatenate(pred, axis=-1)
+            y_pred_test_o = torch.concatenate(pred)
             self.error_test_o = (
                 torch.linalg.norm(y_pred_test_o - y_test_o)
                 / torch.linalg.norm(y_test_o)
@@ -543,8 +524,8 @@ class Trainor_class:
             
             print("Test error on original output: ", self.error_test_o)
 
-            y_test = self.model.norm_out.default(None, y_test_o)
-            y_pred_test = self.model.norm_out.default(None, y_pred_test_o)
+            y_test = pre_func_out(y_test_o)
+            y_pred_test = pre_func_out(y_pred_test_o)
             self.error_test = (
                 torch.linalg.norm(y_pred_test - y_test) / torch.linalg.norm(y_test) * 100
             )
@@ -618,7 +599,7 @@ class Trainor_class:
             dill.dump(obj, f)
         print(f"Object saved in {filename}")
 
-    def load_model(self, filename=None, erase=False, path=None, orig_model_cls=None, **fn_kwargs):
+    def load_model(self, filename=None, erase=False, path=None, orig_model_cls=None, device="cpu",**fn_kwargs):
         """NOTE: fn_kwargs defines the functions of the model
         (e.g. final_activation, inner activation), if
         needed to be saved/loaded on different devices/OS.
@@ -638,29 +619,12 @@ class Trainor_class:
             else:
                 orig_model_cls = orig_model_cls
             kwargs = self.all_kwargs["kwargs"]
-            self.call_map_axis = self.all_kwargs["call_map_axis"]
-            self.call_map_count = self.all_kwargs["call_map_count"]
-            self.params_in = self.all_kwargs["params_in"]
-            self.params_out = self.all_kwargs["params_out"]
-            self.norm_in = self.all_kwargs["norm_in"]
-            self.norm_out = self.all_kwargs["norm_out"]
-            try:
-                self.methods_map = self.all_kwargs["methods_map"]
-                self.methods_norm_in = self.all_kwargs["methods_norm_in"]
-                self.methods_norm_out = self.all_kwargs["methods_norm_out"]
-            except:
-                self.methods_map = ["encode", "decode"]
-                self.methods_norm_in = ["encode"]
-                self.methods_norm_out = ["decode"]
 
             kwargs.update(fn_kwargs)
-
-            model_cls = vmap_wrap(orig_model_cls, self.call_map_axis, self.call_map_count, self.methods_map)
-            model_cls = norm_wrap(model_cls, None, self.norm_in, self.params_in, None, self.norm_out, self.params_out, self.methods_norm_in, self.methods_norm_out)
             
-            model = model_cls(**kwargs)
+            model = orig_model_cls(**kwargs)
             model.load_state_dict(save_dict["model_state_dict"])
-            self.model = model
+            self.model = model.to(device)
             attributes = save_dict["attr"]
 
             for key in attributes:
@@ -671,7 +635,7 @@ class Trainor_class:
 
 class AE_Trainor_class(Trainor_class):
     def __init__(self, *args, **kwargs):
-        super().__init__(*args, methods_map=["encode", "decode"], methods_norm_in=["encode"], methods_norm_out=["decode"], **kwargs)
+        super().__init__(*args, **kwargs)
 
     def fit(self, *args, training_kwargs,  **kwargs):
         if "pre_func_inp" not in kwargs:
@@ -830,11 +794,11 @@ class RRAE_Trainor_class(AE_Trainor_class):
             self.batch_size = 16  # default value
 
         if ft_kwargs:
-            if "get_basis" in ft_kwargs:
-                get_basis = ft_kwargs["get_basis"]
-                ft_kwargs.pop("get_basis")
+            if "get_basis_bool" in ft_kwargs:
+                get_basis_bool = ft_kwargs["get_basis_bool"]
+                ft_kwargs.pop("get_basis_bool")
             else:
-                get_basis = True
+                get_basis_bool = True
             
             if "ft_end_type" in ft_kwargs:
                 ft_end_type = ft_kwargs["ft_end_type"]
@@ -849,8 +813,16 @@ class RRAE_Trainor_class(AE_Trainor_class):
                 ft_end_type = "concat"
                 basis_call_kwargs = {}
 
+            device = ft_kwargs.get("device", "cpu")
+            
+            if "AE_func" in ft_kwargs:
+                AE_func = ft_kwargs["AE_func"]
+                ft_kwargs.pop("AE_func")
+            else:
+                AE_func = lambda m:m
+                
             ft_model, ft_track_params = self.fine_tune_basis(
-                None, args=args, kwargs=ft_kwargs, get_basis=get_basis, end_type=ft_end_type, basis_call_kwargs=basis_call_kwargs
+                None, args=args, kwargs=ft_kwargs, get_basis_bool=get_basis_bool, end_type=ft_end_type, basis_call_kwargs=basis_call_kwargs, device=device, AE_func=AE_func
             )  # fine tune basis
             self.ft_track_params = ft_track_params
         else:
@@ -858,7 +830,7 @@ class RRAE_Trainor_class(AE_Trainor_class):
             ft_track_params = {}
         return model, track_params, ft_model, ft_track_params
 
-    def fine_tune_basis(self, basis=None, get_basis=True, end_type="concat", basis_call_kwargs={}, *, args, kwargs):
+    def fine_tune_basis(self, basis=None, get_basis_bool=True, end_type="concat", basis_call_kwargs={}, device="cpu", AE_func=lambda m:m, *, args, kwargs):
 
         if "loss" in kwargs:
             norm_loss_ = kwargs["loss"]
@@ -869,45 +841,29 @@ class RRAE_Trainor_class(AE_Trainor_class):
             )
 
         if (basis is None):
-            with torch.no_grad():
-                if get_basis:
-                    inp = args[0] if len(args) > 0 else kwargs["input"]
-
-                    if "basis_batch_size" in kwargs:
-                        basis_batch_size = kwargs["basis_batch_size"]
-                        kwargs.pop("basis_batch_size")
-                    else:
-                        basis_batch_size = self.batch_size
+            inp = args[0] if len(args) > 0 else kwargs["input"]
 
-                    basis_kwargs = basis_call_kwargs | self.track_params
-
-                    inpT = inp.permute(*range(inp.ndim - 1, -1, -1))
-                    dataset = TensorDataset(inpT)
-                    dataloader = DataLoader(dataset, batch_size=basis_batch_size, shuffle=False)
-
-                    all_bases = []
-                    
-                    for inp_b in dataloader:
-                        inp_bT = inp_b[0].permute(*range(inp_b[0].ndim - 1, -1, -1))
-                        all_bases.append(self.model.latent(
-                            self.pre_func_inp(inp_bT), get_basis_coeffs=True, **basis_kwargs
-                        )[0]
-                        )
-                    if end_type == "concat":
-                        all_bases = torch.concatenate(all_bases, axis=1)
-                        print(all_bases.shape)
-                        basis = torch.linalg.svd(all_bases, full_matrices=False)[0]
-                        self.basis = basis[:, : self.track_params["k_max"]]
-                    else:
-                        self.basis = all_bases
-                else:
-                    bas = self.model.latent(self.pre_func_inp(inp[..., 0:1]), get_basis_coeffs=True, **self.track_params)[0]
-                    self.basis = torch.eye(bas.shape[0])
+            if "basis_batch_size" in kwargs:
+                basis_batch_size = kwargs["basis_batch_size"]
+                kwargs.pop("basis_batch_size")
+            else:
+                basis_batch_size = self.batch_size
+
+            model = self.model.to(device)
+            k_max = self.track_params["k_max"]
+            if isinstance(AE_func, list):
+                bases = []
+                for func in AE_func:
+                    bases.append(get_basis(get_basis_bool, model, k_max, basis_batch_size, inp, end_type, device, basis_call_kwargs, self.pre_func_inp, func))
+                self.basis = bases
+            else:
+                self.basis = get_basis(get_basis_bool, model, k_max, basis_batch_size, inp, end_type, device, basis_call_kwargs, self.pre_func_inp, AE_func)
+            
         else:
             self.basis = basis
 
         def loss_fun(model, input, out, idx, epsilon, basis):
-            pred = model(input, epsilon=epsilon, apply_basis=basis, keep_normalized=True)
+            pred = model(input, epsilon=epsilon, apply_basis=basis)
             aux = {"loss": norm_loss_(pred, out)}
             return norm_loss_(pred, out), (aux, {})
 
@@ -919,7 +875,7 @@ class RRAE_Trainor_class(AE_Trainor_class):
 
         kwargs.setdefault("loss_kwargs", {}).update({"basis": self.basis})
         
-        fix_comp = lambda model: model._encode.parameters()
+        fix_comp = lambda model: AE_func(model)._encode.parameters()
         print("Fine tuning the basis ...")
         return super().fit(*args, fix_comp=fix_comp, training_kwargs=kwargs)
 
@@ -932,10 +888,10 @@ class RRAE_Trainor_class(AE_Trainor_class):
         batch_size=None,
         pre_func_inp=lambda x: x,
         pre_func_out=lambda x: x,
-        call_func=None,
+        device="cpu",
     ):
 
-        call_func = lambda x: self.model(pre_func_inp(x), apply_basis=self.basis, epsilon=None)
+        call_func = lambda x: self.model(pre_func_inp(x.to(device)), apply_basis=self.basis.to(device), epsilon=None).to("cpu")
         res = super().evaluate(
             x_train_o,
             y_train_o,
@@ -945,6 +901,7 @@ class RRAE_Trainor_class(AE_Trainor_class):
             call_func=call_func,
             pre_func_inp=pre_func_inp,
             pre_func_out=pre_func_out,
+            device=device,
         )
         return res
 

{rraestorch-0.1.5.dist-info → rraestorch-0.1.7.dist-info}/METADATA

@@ -1,11 +1,10 @@
 Metadata-Version: 2.4
 Name: RRAEsTorch
-Version: 0.1.5
+Version: 0.1.7
 Summary:  A repo for RRAEs in PyTorch.
 Author-email: Jad Mounayer <jad.mounayer@outlook.com>
 License: MIT
 License-File: LICENSE
-License-File: LICENSE copy
 Requires-Python: >=3.10
 Requires-Dist: dill
 Requires-Dist: jaxtyping

rraestorch-0.1.7.dist-info/RECORD

@@ -0,0 +1,22 @@
+RRAEsTorch/__init__.py,sha256=f234R6usRCqIgmBmiXyZNIHa7VrDe5E-KZO0Y6Ek5AQ,33
+RRAEsTorch/config.py,sha256=bQPwc_2KTvhglH_WIRSb5_6CpUQQj9AGpfqBp8_kuys,2931
+RRAEsTorch/AE_base/AE_base.py,sha256=Eeo_I7p5P-357rnOmCuFxosJgmBg4KPyMA8n70sTV7U,3368
+RRAEsTorch/AE_base/__init__.py,sha256=95YfMgEWzIFAkm--Ci-a9YPSGfCs2PDAK2sbfScT7oo,24
+RRAEsTorch/AE_classes/AE_classes.py,sha256=iRIcA7iTTDO9Ji90ZKDUmizBswqKRIbnLcwm1hTbVnY,18096
+RRAEsTorch/AE_classes/__init__.py,sha256=inM2_YPJG8T-lwx-CUg-zL2EMltmROQAlNZeZmnvVGA,27
+RRAEsTorch/tests/test_AE_classes_CNN.py,sha256=nxWdxWhpoVbYxMsr-U9mmd7vZ8s0pBaLsDvCsti2r-I,2913
+RRAEsTorch/tests/test_AE_classes_MLP.py,sha256=Aa8PhhsWblw8tJATVEfw9DLFA5vK-0lSzG6BySwmTiM,2413
+RRAEsTorch/tests/test_fitting_CNN.py,sha256=VlWXxFMDul1MTMt1TciezYZpKifLTl1YsIFUAlAxc1w,2796
+RRAEsTorch/tests/test_fitting_MLP.py,sha256=FVw6ObzOZJG5CntHoW45A1me5QFgIBI64rYPxjh1ujM,3288
+RRAEsTorch/tests/test_mains.py,sha256=ivTXP7NypSlgmB9VR5g0yq5VEuPZJGOibDqBMjOxHow,1021
+RRAEsTorch/tests/test_save.py,sha256=KdXwoF3ao7UVr8tEJ7sEHqNtUqIYmuoVEAvt4FigXIs,1982
+RRAEsTorch/tests/test_stable_SVD.py,sha256=OimHPqw4f22qndyRzwJfNvTzzjP2CM-yHtfXCqkMBuA,1230
+RRAEsTorch/trackers/__init__.py,sha256=3c9qcUMZiUfVr93rxFp6l11lIDthyK3PCY_-P-sNX3I,25
+RRAEsTorch/trackers/trackers.py,sha256=Pn1ejMxMjAtvgDazFFwa3qiZhogG5GtXj4UIIFiBpuY,9127
+RRAEsTorch/training_classes/__init__.py,sha256=K_Id4yhw640jp2JN15-0E4wJi4sPadi1fFRgovMV3kw,101
+RRAEsTorch/training_classes/training_classes.py,sha256=tZedyKvi3f-phiByIG820Z_oMsG5vk9xQK_OS_04jIM,34014
+RRAEsTorch/utilities/__init__.py,sha256=NtlizCcRW4qcsULXxWfjPk265rLJst0-GqWLRah2yDY,26
+rraestorch-0.1.7.dist-info/METADATA,sha256=fdYJFLO7CHpG2IxfaC1zTFwiA6Lf2cyPqHiueU0lvs0,3028
+rraestorch-0.1.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+rraestorch-0.1.7.dist-info/licenses/LICENSE,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
+rraestorch-0.1.7.dist-info/RECORD,,

RRAEsTorch/tests/test_wrappers.py

@@ -1,56 +0,0 @@
-from RRAEsTorch.wrappers import vmap_wrap, norm_wrap
-from torchvision.ops import MLP
-import pytest
-import numpy as np
-import math
-import numpy.random as random
-import torch
-
-def test_vmap_wrapper():
-    # Usually MLP only accepts a vector, here we give
-    # a tensor and vectorize over the last axis twice
-    data = random.normal(size=(50, 60, 600))
-    data = torch.tensor(data, dtype=torch.float32)
-    
-    model_cls = vmap_wrap(MLP, -1, 2)
-    model = model_cls(50, [64, 100])
-    try:
-        model(data)
-    except ValueError:
-        pytest.fail("Vmap wrapper is not working properly.")
-
-def test_norm_wrapper():
-    # Testing the keep_normalized kwarg
-    data = random.normal(size=(50,))
-    data = torch.tensor(data, dtype=torch.float32)
-    model_cls = norm_wrap(MLP, data, "minmax", None, data, "minmax", None)
-    model = model_cls(50, [64, 100])
-    try:
-        assert not torch.allclose(model(data), model(data, keep_normalized=True))
-    except AssertionError:
-        pytest.fail("The keep_normalized kwarg for norm wrapper is not behaving as expected.")
-
-    # Testing minmax with knwon mins and maxs
-    data = np.linspace(-1, 1, 100)
-    data = torch.tensor(data, dtype=torch.float32)
-    model_cls = norm_wrap(MLP, data, "minmax", None, data, "minmax", None)
-    model = model_cls(50, [64, 100])
-    try:
-        assert 0.55 == model.norm_in.default(None, 0.1)
-        assert -0.8 == model.inv_norm_out.default(None, 0.1)
-    except AssertionError:
-        pytest.fail("Something wrong with minmax wrapper.")
-
-    # Testing meanstd with knwon mean and std
-    data = random.normal(size=(50,))
-    data = (data-np.mean(data))/np.std(data)
-    data = data*2.0 + 1.0  # mean of 1 and std of 2
-    data = torch.tensor(data, dtype=torch.float32)
-
-    model_cls = norm_wrap(MLP, data, "meanstd", None, data, "meanstd", None)
-    model = model_cls(50, [64, 100])
-    try:
-        assert math.isclose(2, model.norm_in.default(None, 5), rel_tol=1e-1, abs_tol=1e-1)
-        assert math.isclose(7, model.inv_norm_out.default(None, 3), rel_tol=1e-1, abs_tol=1e-1)
-    except AssertionError:
-        pytest.fail("Something wrong with norm wrapper.")