rxnn 0.2.42__tar.gz → 0.2.44__tar.gz

{rxnn-0.2.42 → rxnn-0.2.44}/PKG-INFO

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

{rxnn-0.2.42 → rxnn-0.2.44}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.42"
+version = "0.2.44"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.2.42 → rxnn-0.2.44}/src/rxnn/memory/stm.py

@@ -62,6 +62,9 @@ class ShortTermMemory(nn.Module):
     def reset(self, init_type: str = None):
         self.memory = self._init_tensor(init_type).to(self.memory.device)
 
+    def clone_detach_reset(self):
+        self.memory = self.memory.detach().clone()
+
     def resize(self, new_stm_size: int, init_type: str = None):
         self.stm_size = new_stm_size
         delattr(self, 'memory')

{rxnn-0.2.42 → rxnn-0.2.44}/src/rxnn/rxt/models.py

@@ -301,6 +301,9 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
     def reset_memory(self, init_type: str = None):
         self.model.stm.reset(init_type)
 
+    def clone_reset_memory(self):
+        self.model.stm.clone_detach_reset()
+
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         return self.model(x, attention_mask=attention_mask)
 

{rxnn-0.2.42 → rxnn-0.2.44}/src/rxnn/training/models.py

@@ -146,6 +146,9 @@ class MrlActorModel(nn.Module):
     def reset_memory(self):
         self.memory_attention.reset_memory()
 
+    def clone_reset_memory(self):
+        self.memory_attention.clone_reset_memory()
+
     def memory_parameters(self) -> list[nn.Parameter]:
         return list(set(
             self.encoder.memory_parameters() +
@@ -168,6 +171,9 @@ class MrlActorModel(nn.Module):
             self.decoder.not_memory_parameters()
         ))
 
+    def embedding_parameters(self) -> Iterator[nn.Parameter]:
+        return self.encoder.model.embedding.parameters()
+
     def unique_parameters(self, with_embedding: bool = True):
         if with_embedding:
             return list(set(

{rxnn-0.2.42 → rxnn-0.2.44}/src/rxnn/training/mrl.py

@@ -225,7 +225,7 @@ class MRLTrainer:
     ) -> tuple[torch.optim.Optimizer, torch.optim.Optimizer]:
         if memory_lr is not None:
             optimizer = torch.optim.AdamW([
-                {'params': self.actor.encoder.embedding.parameters(), 'lr': embedding_lr},
+                {'params': self.actor.embedding_parameters(), 'lr': embedding_lr},
                 {'params': self.actor.not_memory_parameters(), 'lr': lr},
                 {'params': self.actor.memory_parameters(), 'lr': memory_lr},
             ],
@@ -233,7 +233,7 @@ class MRLTrainer:
             )
         else:
             optimizer = torch.optim.AdamW([
-                {'params': self.actor.encoder.embedding.parameters(), 'lr': embedding_lr},
+                {'params': self.actor.embedding_parameters(), 'lr': embedding_lr},
                 {'params': self.actor.unique_parameters(with_embedding=False), 'lr': lr},
             ],
                 weight_decay=weight_decay,
@@ -646,6 +646,8 @@ class MRLTrainer:
                 step_critic_values = episode_critic_values[step_idx]
                 step_advantages = episode_advantages[step_idx]
 
+                self.actor.clone_reset_memory()
+
                 # 7. In memory aware critic version, encode and update STM before critic update, to include its gradients in critic loss too
                 if self.memory_aware_critic:
                     self.encode_and_update_stm(query, answer)
@@ -929,7 +931,7 @@ class MRLTrainer:
 
         if mode == 'update':
             params = [
-                {'params': self.actor.encoder.embedding.parameters(), 'lr': embedding_lr},
+                {'params': self.actor.embedding_parameters(), 'lr': embedding_lr},
                 {'params': self.actor.encoder.not_memory_parameters(), 'lr': model_lr},
                 {'params': self.actor.encoder.memory_parameters(), 'lr': memory_lr},
                 {'params': self.actor.memory_attention_parameters(), 'lr': memory_lr},
@@ -938,7 +940,7 @@ class MRLTrainer:
             ]
         elif mode == 'fetch':
             params = [
-                {'params': self.actor.encoder.embedding.parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.embedding_parameters(), 'lr': unfreeze_lr},
                 {'params': self.actor.encoder.not_memory_parameters(), 'lr': unfreeze_lr},
                 {'params': self.actor.encoder.memory_parameters(), 'lr': unfreeze_lr},
                 {'params': self.actor.memory_attention_parameters(), 'lr': unfreeze_lr},
@@ -947,7 +949,7 @@ class MRLTrainer:
             ]
         elif mode == 'joint':
             params = [
-                {'params': self.actor.encoder.embedding.parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.embedding_parameters(), 'lr': unfreeze_lr},
                 {'params': self.actor.encoder.not_memory_parameters(), 'lr': unfreeze_lr},
                 {'params': self.actor.encoder.memory_parameters(), 'lr': memory_lr},
                 {'params': self.actor.memory_attention_parameters(), 'lr': memory_lr},
@@ -956,7 +958,7 @@ class MRLTrainer:
             ]
         else:
             params = [
-                {'params': self.actor.encoder.embedding.parameters(), 'lr': embedding_lr},
+                {'params': self.actor.embedding_parameters(), 'lr': embedding_lr},
                 {'params': self.actor.encoder.not_memory_parameters(), 'lr': model_lr},
                 {'params': self.actor.encoder.memory_parameters(), 'lr': memory_lr},
                 {'params': self.actor.memory_attention_parameters(), 'lr': memory_lr},