rxnn 0.2.70__tar.gz → 0.2.71__tar.gz

{rxnn-0.2.70 → rxnn-0.2.71}/PKG-INFO

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

{rxnn-0.2.70 → rxnn-0.2.71}/pyproject.toml

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

{rxnn-0.2.70 → rxnn-0.2.71}/src/rxnn/training/mrl.py

@@ -607,21 +607,38 @@ class MRLTrainer:
         print(f"Encoder grad norm - total: {encoder_total:.6f}, mean: {encoder_mean:.6f}")
         print(f"Decoder grad norm - total: {decoder_total:.6f}, mean: {decoder_mean:.6f}")
         print(f"Memory attention grad norm - total: {mem_att_total:.6f}, mean: {mem_att_mean:.6f}")
-        # decoder's cross att
-        dec_x_att_norms = [get_gradient_norms(layer.memory_cross_attention)[1] for layer in self.actor.decoder.model.layers]
-        print(f"Decoder cross-att mean norm: {(sum(dec_x_att_norms) / len(dec_x_att_norms)):.6f}, all: {dec_x_att_norms}")
 
+        dec_x_att_norms = [get_gradient_norms(layer.memory_cross_attention)[1] for layer in self.actor.decoder.model.layers]
         mem_att_norms = [get_gradient_norms(layer)[1] for layer in self.actor.memory_attention.model.attention_layers]
-        print(f"Memory attention layers mean norm: {(sum(mem_att_norms) / len(mem_att_norms)):.6f}, all: {mem_att_norms}")
-
         enc_ff_norms = [get_gradient_norms(layer.ff)[1] for layer in self.actor.encoder.model.layers]
-        print(f"Encoder ff mean norm: {(sum(enc_ff_norms) / len(enc_ff_norms)):.6f}, all: {enc_ff_norms}")
-
         enc_self_att_norms = [get_gradient_norms(layer.attention)[1] for layer in self.actor.encoder.model.layers]
-        print(f"Encoder self-att mean norm: {(sum(enc_self_att_norms) / len(enc_self_att_norms)):.6f}, all: {enc_self_att_norms}")
+        enc_x_att_norms = [get_gradient_norms(layer.memory_cross_attention)[1] for layer in
+                         self.actor.encoder.model.layers]
+
+        calc_mean = lambda x: sum(x) / len(x)
+
+        dec_x_att_norms_mean = calc_mean(dec_x_att_norms)
+        mem_att_norms_mean = calc_mean(mem_att_norms)
+        enc_ff_norms_mean = calc_mean(enc_ff_norms)
+        enc_self_att_norms_mean = calc_mean(enc_self_att_norms)
+        enc_x_att_norms_mean = calc_mean(enc_x_att_norms)
+
+        print(f"Decoder cross-att mean norm: {dec_x_att_norms_mean:.6f}, all: {dec_x_att_norms}")
+        print(f"Memory attention layers mean norm: {mem_att_norms_mean:.6f}, all: {mem_att_norms}")
+        print(f"Encoder ff mean norm: {enc_ff_norms_mean:.6f}, all: {enc_ff_norms}")
+        print(f"Encoder self-att mean norm: {enc_self_att_norms_mean:.6f}, all: {enc_self_att_norms}")
+        print(f"Encoder cross-att mean norm: {enc_x_att_norms_mean:.6f}, all: {enc_x_att_norms}")
+
+        if self.writer is not None:
+            self.writer.add_scalar('Gradient/encoder', encoder_mean, self.global_step['train'])
+            self.writer.add_scalar('Gradient/decoder', decoder_mean, self.global_step['train'])
+            self.writer.add_scalar('Gradient/mem-att', mem_att_mean, self.global_step['train'])
+            self.writer.add_scalar('Gradient/decoder x-att', dec_x_att_norms_mean, self.global_step['train'])
+            self.writer.add_scalar('Gradient/mem-att layers', mem_att_norms_mean, self.global_step['train'])
+            self.writer.add_scalar('Gradient/encoder ff', enc_ff_norms_mean, self.global_step['train'])
+            self.writer.add_scalar('Gradient/encoder self-att', enc_self_att_norms_mean, self.global_step['train'])
+            self.writer.add_scalar('Gradient/encoder x-att', enc_x_att_norms_mean, self.global_step['train'])
 
-        enc_att_norms = [get_gradient_norms(layer.memory_cross_attention)[1] for layer in self.actor.encoder.model.layers]
-        print(f"Encoder cross-att mean norm: {(sum(enc_att_norms) / len(enc_att_norms)):.6f}, all: {enc_att_norms}")
 
     def update_actor(self, state: tuple[TokenizedDict, TokenizedDict, TokenizedDict], action: TokenizedDict,
                      advantages: torch.Tensor, old_log_probs: torch.Tensor, epoch: int) -> float: