rxnn 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl

rxnn/memory/attention.py

@@ -36,8 +36,7 @@ class StmMemoryAttention(nn.Module):
             normalized_layer_stm = self.memory_norm_layers[i](layer_stm)
             new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=mask)
             # self.stm.update_layer(i, new_layer_stm + layer_stm)
-            final_layer_stm = new_layer_stm + layer_stm # residual
-            self.stm.update_layer(i, final_layer_stm)
-        # self.stm.update_all(new_stm)
+            new_stm[i] = new_layer_stm + layer_stm # residual
+        self.stm.update_all(new_stm)
         return self.stm.memory
 

rxnn/memory/stm.py

@@ -39,10 +39,12 @@ class ShortTermMemory(nn.Module):
         return self.memory[layer]
 
     def update_layer(self, layer: int, new_stm: torch.Tensor):
+        self.memory = self.memory.clone()
         self.memory[layer] = new_stm
 
     def update_all(self, new_stm: torch.Tensor):
-        self.memory.copy_(new_stm)
+        self.memory = new_stm
+        # self.memory.copy_(new_stm)
 
     def make_trainable(self):
         if not self.is_trainable:
@@ -59,7 +61,7 @@ class ShortTermMemory(nn.Module):
             self.register_buffer('memory', trained_stm)
 
     def reset(self, init_type: str = None):
-        self.memory.copy_(self._init_tensor(init_type))
+        self.memory = self._init_tensor(init_type).to(self.memory.device)
 
     def resize(self, new_stm_size: int, init_type: str = None):
         self.stm_size = new_stm_size
@@ -84,8 +86,7 @@ class ShortTermMemory(nn.Module):
         if use_mean_from_batch:
             batch_mean = self.memory.mean(dim=(1, 2, 3), keepdim=True)
             delattr(self, 'memory')
-            self.register_buffer('memory', self._init_tensor())
-            self.memory.copy_(batch_mean)
+            self.register_buffer('memory', batch_mean)
         else:
             delattr(self, 'memory')
             self.register_buffer('memory', self._init_tensor())

rxnn/training/mrl.py

@@ -461,7 +461,6 @@ class MRLTrainer:
         # 1. Run update separately for episodes in trajectory - we have to reset memory before each episode, and update
         # memory, based on collected episode data
         all_losses = []
-        trajectories_len = len(trajectories)
         for episode_idx, episode in enumerate(trajectories):
             episode_steps = episode['steps']
             should_reset_stm = episode['reset_stm']
@@ -514,14 +513,14 @@ class MRLTrainer:
 
                 # 9. Update the model in AMP or regular mode
                 if self.use_amp:
-                    self.scaler.scale(policy_loss).backward()
+                    self.scaler.scale(policy_loss).backward(retain_graph=True)
                     self.scaler.unscale_(self.optimizer)
                     torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                                    error_if_nonfinite=False)
                     self.scaler.step(self.optimizer)
                     self.scaler.update()
                 else:
-                    policy_loss.backward()
+                    policy_loss.backward(retain_graph=True)
                     torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                                    error_if_nonfinite=False)
                     self.optimizer.step()

rxnn/training/rl.py

@@ -43,6 +43,8 @@ class PPOAlgorithm(RlAlgorithm):
         # b) Calculate ratio
         ratio = (new_log_probs - old_log_probs).exp()
 
+        advantages = advantages.unsqueeze(-1)
+
         # c) Clipped surrogate loss
         surr1 = ratio * advantages
         surr2 = torch.clamp(ratio, 1.0 - self.clip_eps, 1.0 + self.clip_eps) * advantages

{rxnn-0.2.11.dist-info → rxnn-0.2.13.dist-info}/METADATA

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

{rxnn-0.2.11.dist-info → rxnn-0.2.13.dist-info}/RECORD

@@ -5,9 +5,9 @@ rxnn/experimental/attention.py,sha256=46qwZLJuZMpIBrZ-r9DaQEPPmmZkO464C3Tkm_Mq-c
 rxnn/experimental/models.py,sha256=foBo0n0ufvBnfIdJomiEg3CuSOiWSt-q5ako7vzYxx4,4888
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/memory/attention.py,sha256=eCsTjJEQVurBjsIlJyk1cDvOGU2YMbYMKiBVOb3mfKg,1874
+rxnn/memory/attention.py,sha256=p-r8DK3iVhNn-JAESVzIXDCG8gk1R_-x5xHclZ5jgb0,1813
 rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
-rxnn/memory/stm.py,sha256=DPkK1q1SLRw3HWM0dcvkn4XvIrfwUK47h4KmvFVWljc,3847
+rxnn/memory/stm.py,sha256=IH_3INw7FdI013t56ui3Zq9GPUq-k3HeZGjx6BerS4g,3888
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/rxt/models.py,sha256=qlYgU002VE21ZOlcxEM9iv9tAvsbe4mngcMI2sw3j9k,12078
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -16,9 +16,9 @@ rxnn/training/bml.py,sha256=FJszaQXOLx2ZHBa1CQpyMrG8i4Kj14E-gzDAEK_Ei5k,17272
 rxnn/training/callbacks.py,sha256=-N0MQPpZQaUWCINdTOsjul4bDGbGr2JgQBqOIXBLS6o,35053
 rxnn/training/dataset.py,sha256=i8EdSJnoPbkuDSdqoYDj-Aig5Se_uPY4lulkD2bdOrs,50331
 rxnn/training/models.py,sha256=renPa5YH443XNTMFI-YTCwi5vNp3QzwF5UXedNd5hDk,5187
-rxnn/training/mrl.py,sha256=mCsg50bX0iqPozvvQB6CeZ0FYEfuj9ln1p-4IaZBryo,39338
+rxnn/training/mrl.py,sha256=53uOwotmgwKeceMYA6qXQbQMZmggXt_5hq08X-YwrEY,39327
 rxnn/training/reward.py,sha256=C0ToTz-u-L-qyBd2yJ1HlvVPS110OChYj9ZhD6iSSMU,5654
-rxnn/training/rl.py,sha256=T69gLwDlvMMyLuRaJSRmwzO0Mcu0uLXwhAiBB58VK-Y,2663
+rxnn/training/rl.py,sha256=s6wPbg0X6y-RX9-5ctZIDpdJPfExI9DzWUy-TvAiiow,2710
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
@@ -32,7 +32,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.11.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.11.dist-info/METADATA,sha256=Latj91qV2ruPOzt-bdaqfGfu0WF42e5qPz6MCBVVtTo,25960
-rxnn-0.2.11.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.11.dist-info/RECORD,,
+rxnn-0.2.13.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.13.dist-info/METADATA,sha256=BOn4qig3IKpYiG0NEWHiF_5NWsWboBqVNeGb2-mYesU,25960
+rxnn-0.2.13.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.13.dist-info/RECORD,,