rxnn 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

rxnn/memory/stm.py

@@ -5,7 +5,7 @@ class ShortTermMemory(nn.Module):
     """Short-term memory module for the Attention-based Memory System"""
 
     def __init__(self, num_layers: int, embed_dim: int, stm_size: int, init_type: str = 'normal',
-                 is_trainable: bool = False, legacy_init: bool = True, *args, **kwargs):
+                 is_trainable: bool = False, *args, **kwargs):
         super(ShortTermMemory, self).__init__(*args, **kwargs)
         self.num_layers = num_layers
         self.embed_dim = embed_dim
@@ -14,10 +14,6 @@ class ShortTermMemory(nn.Module):
         self.is_trainable = is_trainable
         assert init_type in ['normal', 'standard', 'uniform', 'ones', 'zeros'], \
             'STM init type must be one of "normal", "standard", "uniform", "ones", "zeros"'
-
-        # Legacy init - temporary option to load old models with not-batched STM (they will be loaded, updated and then the option will be removed)
-        self.legacy_init = legacy_init
-
         self.init_type = init_type
         stm = self._init_tensor()
         if self.is_trainable:
@@ -27,8 +23,7 @@ class ShortTermMemory(nn.Module):
 
     def _init_tensor(self, init_type: str = None):
         init_type = init_type or self.init_type
-        stm_shape = (self.num_layers, self.stm_size, self.embed_dim) \
-            if self.legacy_init else (self.num_layers, self.batch_size, self.stm_size, self.embed_dim)
+        stm_shape = (self.num_layers, self.batch_size, self.stm_size, self.embed_dim)
         if init_type == 'normal':
             return torch.normal(0, 0.02, stm_shape)
         elif init_type == 'standard':
@@ -40,12 +35,8 @@ class ShortTermMemory(nn.Module):
         else:
             return torch.zeros(*stm_shape)
 
-    def reset_legacy_(self):
-        self.legacy_init = False
-        self.memory = self._init_tensor()
-
     def forward(self, layer: int) -> torch.Tensor:
-        return self.memory[layer].unsqueeze(0) if self.legacy_init else self.memory[layer]
+        return self.memory[layer]
 
     def update_layer(self, layer: int, new_stm: torch.Tensor):
         self.memory[layer] = new_stm

rxnn/training/callbacks.py

@@ -557,7 +557,7 @@ class MrlPrintCallback(MrlTrainerCallback):
     def on_epoch_start(self, actor: nn.Module, epoch: int, stage_epochs: int, curriculum_config: dict,
                        global_epoch: int, global_epochs: int) -> None:
         print(
-            f'Starting epoch {epoch}/{stage_epochs} (stage) | {global_epoch}/{global_epochs} (global) for {curriculum_config['steps']} steps in {curriculum_config['strategy']} strategy.')
+            f'Starting epoch {epoch}/{stage_epochs} (stage) | {global_epoch}/{global_epochs} (global) for {curriculum_config["steps"]} steps in {curriculum_config["strategy"]} strategy.')
 
     def on_epoch_end(self, actor: nn.Module, epoch: int, stage_epochs: int, policy_loss: float,
                      critic_loss: float, global_epoch: int, global_epochs: int) -> None:
@@ -580,7 +580,7 @@ class MrlPrintCallback(MrlTrainerCallback):
         print(f'Epoch {epoch} | Step {step} - updated policy loss {critic_loss}')
 
     def on_training_end(self, actor: nn.Module, critic: nn.Module, curriculum_config: dict) -> None:
-        print(f'Finished training for {curriculum_config['steps']} steps in {curriculum_config['strategy']} strategy.')
+        print(f'Finished training for {curriculum_config["steps"]} steps in {curriculum_config["strategy"]} strategy.')
 
     def on_eval_end(self, actor: nn.Module, critic: nn.Module, epoch: int, eval_mean_reward: float) -> None:
         print(f'Eval epoch {epoch} - mean reward {eval_mean_reward}')

{rxnn-0.2.1.dist-info → rxnn-0.2.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.1
+Version: 0.2.3
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.2.1.dist-info → rxnn-0.2.3.dist-info}/RECORD

@@ -7,13 +7,13 @@ rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/attention.py,sha256=p-r8DK3iVhNn-JAESVzIXDCG8gk1R_-x5xHclZ5jgb0,1813
 rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
-rxnn/memory/stm.py,sha256=tEAC5c4QPl-ykNm-k2DlcP6Q2NQKGXMnpKLCx92lI30,4083
+rxnn/memory/stm.py,sha256=eSMK5KdupWNf56FcDYprHnjA51EeYBzSKza7tiZxKSc,3618
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/rxt/models.py,sha256=zNrf6mn-s2vJyauHwNgYm_e-gFI1clmXp_JyCKGQD3E,12083
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=_xik1GXE4RJ_nxwqLQ1ccXA5pRtBCi-jL-jeRFBdHBU,11851
 rxnn/training/bml.py,sha256=FJszaQXOLx2ZHBa1CQpyMrG8i4Kj14E-gzDAEK_Ei5k,17272
-rxnn/training/callbacks.py,sha256=aqi8CfXUWnjMDbELYC5BPBbYyq0YiMicyVaTIr778DY,35053
+rxnn/training/callbacks.py,sha256=o68IPFJyWM1CGooPRDNU9DfNcy4H_o0PcKDTn_ZLnKA,35053
 rxnn/training/dataset.py,sha256=XeRzo0KUYyQ43XjZ3o6Jban9ePIRtpHsqUmeKAQPRQk,50305
 rxnn/training/models.py,sha256=renPa5YH443XNTMFI-YTCwi5vNp3QzwF5UXedNd5hDk,5187
 rxnn/training/mrl.py,sha256=KcGvBWlBcFJ5GSwd4lx3pUXKlcyeNgJYPZAk3DRMH48,39179
@@ -32,7 +32,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=2dpUQv88ekZa_CMSPLrXvB6X684wxUE2bDVznsi5ACs,17429
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.1.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.1.dist-info/METADATA,sha256=0NaOW0HNG4nGR3Vchlp-vlzzx7NBSFpAPFYwGzT3qP8,25959
-rxnn-0.2.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.1.dist-info/RECORD,,
+rxnn-0.2.3.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.3.dist-info/METADATA,sha256=qe9tpeNRuSsTeABm7K2hUPyAoRPNt-ZFHl_5Nv_u7Mc,25959
+rxnn-0.2.3.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.3.dist-info/RECORD,,