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

RRAEsTorch/AE_classes/AE_classes.py

@@ -127,7 +127,7 @@ def latent_func_var_strong_RRAE(self, y, k_max=None, epsilon=None, return_dist=F
     if epsilon is not None:
         if len(epsilon.shape) == 4:
             epsilon = epsilon[0, 0] # to allow tpu sharding
-        z = mean + torch.tensor(epsilon, dtype=torch.float32) * std
+        z = mean + epsilon * std
     else:
         z = mean
 
@@ -612,10 +612,6 @@ class VRRAE_CNN1D(CNN1D_Autoencoder):
     typ: int
 
     def __init__(self, channels, input_dim, latent_size, k_max, typ="eye", *, count=1, **kwargs):
-        v_Linear = vmap_wrap(Linear, -1, count=count)
-        self.lin_mean = v_Linear(k_max, k_max,)
-        self.lin_logvar = v_Linear(k_max, k_max)
-        self.typ = typ
         super().__init__(
             channels,
             input_dim,
@@ -623,6 +619,12 @@ class VRRAE_CNN1D(CNN1D_Autoencoder):
             count=count,
             **kwargs,
         )
+        
+        v_Linear = vmap_wrap(Linear, -1, count=count)
+        self.lin_mean = v_Linear(k_max, k_max,)
+        self.lin_logvar = v_Linear(k_max, k_max)
+        self.typ = typ
+        
 
     def _perform_in_latent(self, y, *args, k_max=None, epsilon=None, return_dist=False, return_lat_dist=False, **kwargs):
         return latent_func_var_strong_RRAE(self, y, k_max, epsilon, return_dist, return_lat_dist, **kwargs)

RRAEsTorch/training_classes/training_classes.py

@@ -229,7 +229,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 +302,12 @@ class Trainor_class:
             all_losses = []
 
         
-        
+        model = model.to(device)
+
+        if input_val is not None:
+            dataset_val = TensorDataset(input_val.permute(*range(input_val.ndim - 1, -1, -1)), output_val.permute(*range(output_val.ndim - 1, -1, -1)), torch.arange(0, input_val.shape[-1], 1))
+            dataloader_val = DataLoader(dataset_val, batch_size=input_val.shape[-1], shuffle=False)
+
         # Outler Loop
         for steps, lr, batch_size in zip(step_st, lr_st, batch_size_st):
             try:
@@ -340,6 +346,9 @@ class Trainor_class:
                     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 +360,17 @@ class Trainor_class:
                                                             )
                     
                     if input_val is not None:
+                        val_loss = []
+                        for input_vb, out_vb, idx_b in dataloader_val:
+                            input_vb = input_vb.permute(*range(input_vb.ndim - 1, -1, -1))
+                            out_vb = self.model.norm_out.default(None, pre_func_out(out_vb))    # Pre-process batch out values
+                            out_vb = out_vb.permute(*range(out_vb.ndim - 1, -1, -1))
     
-
-                        idx = np.arange(input_val.shape[-1])
-                        val_loss, _ = loss_fun(
-                            model, input_val, output_val, idx=idx, epsilon=None, **step_kwargs
-                        )
+                            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
@@ -618,7 +632,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.
@@ -660,7 +674,7 @@ class Trainor_class:
             
             model = 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:
@@ -849,8 +863,10 @@ class RRAE_Trainor_class(AE_Trainor_class):
                 ft_end_type = "concat"
                 basis_call_kwargs = {}
 
+            device = ft_kwargs.get("device", "cpu")
+            
             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=get_basis, end_type=ft_end_type, basis_call_kwargs=basis_call_kwargs, device=device
             )  # fine tune basis
             self.ft_track_params = ft_track_params
         else:
@@ -858,7 +874,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=True, end_type="concat", basis_call_kwargs={}, device="cpu", *, args, kwargs):
 
         if "loss" in kwargs:
             norm_loss_ = kwargs["loss"]
@@ -884,24 +900,24 @@ class RRAE_Trainor_class(AE_Trainor_class):
                     inpT = inp.permute(*range(inp.ndim - 1, -1, -1))
                     dataset = TensorDataset(inpT)
                     dataloader = DataLoader(dataset, batch_size=basis_batch_size, shuffle=False)
+                    model = self.model.to(device)
 
                     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]
+                        all_bases.append(model.latent(
+                            self.pre_func_inp(inp_bT.to(device)), get_basis_coeffs=True, **basis_kwargs
+                        )[0].to("cpu")
                         )
                     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]
+                    bas = model.latent(self.pre_func_inp(inp[..., 0:1].to(device)), get_basis_coeffs=True, **self.track_params)[0].to("cpu")
                     self.basis = torch.eye(bas.shape[0])
         else:
             self.basis = basis
@@ -932,10 +948,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 +961,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/utilities/utilities.py

@@ -1510,7 +1510,6 @@ class MLP_with_CNN3D_trans(torch.nn.Module):
         self.out_after_mlp = out_after_mlp
 
     def forward(self, x, *args, **kwargs):
-        print(x.shape)
         x = self.layers[0](x)
         x = torch.reshape(x, (self.out_after_mlp, self.first_D0, self.first_D1, self.first_D2))
         x = self.layers[1](x)

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

@@ -1,11 +1,10 @@
 Metadata-Version: 2.4
 Name: RRAEsTorch
-Version: 0.1.5
+Version: 0.1.6
 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.5.dist-info → rraestorch-0.1.6.dist-info}/RECORD

@@ -2,7 +2,7 @@ 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=HUiCH5iXhdMSZIQiulIktIwCpDjfPAXPrLiMpUYHg4M,18078
+RRAEsTorch/AE_classes/AE_classes.py,sha256=oDpDzQasPbtK2L9vDLiG4VQdKH02VRCagOYT1-FAldo,18063
 RRAEsTorch/AE_classes/__init__.py,sha256=inM2_YPJG8T-lwx-CUg-zL2EMltmROQAlNZeZmnvVGA,27
 RRAEsTorch/tests/test_AE_classes_CNN.py,sha256=bEE9JnTo84t9w0a4kw1W74L51eLGjBB8trrlAG938RE,3182
 RRAEsTorch/tests/test_AE_classes_MLP.py,sha256=Cr1_uP7lag6RPQC1UhN2O7RFW4BEx1cd0Z-Y6VgrWRg,2718
@@ -15,13 +15,12 @@ RRAEsTorch/tests/test_wrappers.py,sha256=Ike4IfMUx2Qic3f3_cBikgFPEU1WW5TuH1jT_r2
 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=a7JjhCrH7s7VmVCsvKj768Ciq-tdbh6E_B9aG1kw7vc,36634
+RRAEsTorch/training_classes/training_classes.py,sha256=HU8Ksz1-2WwOMuwyGiWdkQ_vrrgBEwaeQT4avs4jd2E,37870
 RRAEsTorch/utilities/__init__.py,sha256=NtlizCcRW4qcsULXxWfjPk265rLJst0-GqWLRah2yDY,26
-RRAEsTorch/utilities/utilities.py,sha256=FzJWV9oFPF9sL9MC2m7euMqMKxCuLUEukzLfU0cF2to,53396
+RRAEsTorch/utilities/utilities.py,sha256=JfLkAPEC8fzwgM32LEcXVe0tA4C7UBgsrkuh6noUA_4,53372
 RRAEsTorch/wrappers/__init__.py,sha256=txiLh4ylnuvPlapagz7DiAslmjllOzTqwCDL2dFr6dM,44
 RRAEsTorch/wrappers/wrappers.py,sha256=9Rmq2RS_EkZvsg96SKrt1HFIP35sF0xyPI0goV0ujOs,9659
-rraestorch-0.1.5.dist-info/METADATA,sha256=uL69l4DSjgNliJfX3uPyrZqrH2qvqZAu55DE1jfF5KQ,3055
-rraestorch-0.1.5.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-rraestorch-0.1.5.dist-info/licenses/LICENSE,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
-rraestorch-0.1.5.dist-info/licenses/LICENSE copy,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
-rraestorch-0.1.5.dist-info/RECORD,,
+rraestorch-0.1.6.dist-info/METADATA,sha256=BfnB-vhx0m-d79hCd8UgZT0c8GPHFRwVn8x-M8k_h6E,3028
+rraestorch-0.1.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+rraestorch-0.1.6.dist-info/licenses/LICENSE,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
+rraestorch-0.1.6.dist-info/RECORD,,

rraestorch-0.1.5.dist-info/licenses/LICENSE copy

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2026 Jad Mounayer
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.