smaller task train

experiments/meta_task_exp.py

@@ -295,8 +295,7 @@ def flat_for_store(parameters):
 
 
 def train_self_replication(model, optimizer, st_stps) -> dict:
-    for _ in range(st_stps):
-        self_train_loss = model.combined_self_train(optimizer)
+    self_train_loss = model.combined_self_train(optimizer, st_stps)
     # noinspection PyUnboundLocalVariable
     stp_log = dict(Metric='Self Train Loss', Score=self_train_loss.item())
     return stp_log

experiments/meta_task_sanity_exp.py

@@ -68,7 +68,8 @@ if __name__ == '__main__':
         mean_self_tain_loss = []
 
         for batch, (batch_x, batch_y) in tenumerate(dataloader):
-            # self_train_loss, _ = net.self_train(2, save_history=False, learning_rate=0.004)
+            self_train_loss, _ = net.self_train(2, save_history=False, learning_rate=0.004)
+
             for _ in range(2):
                 optimizer.zero_grad()
                 input_data = net.input_weight_matrix()

network.py

@@ -463,19 +463,20 @@ class MetaNet(nn.Module):
     def particles(self):
         return (cell for metalayer in self.all_layers for cell in metalayer.particles)
 
-    def combined_self_train(self, optimizer, reduction='mean'):
-        optimizer.zero_grad()
+    def combined_self_train(self, optimizer, n_st_steps, reduction='mean'):
+
         losses = []
-        n = 10
         for particle in self.particles:
+            for _ in range(n_st_steps):
+                optimizer.zero_grad()
                 # Intergrate optimizer and backward function
                 input_data = particle.input_weight_matrix()
                 target_data = particle.create_target_weights(input_data)
                 output = particle(input_data)
                 losses.append(F.mse_loss(output, target_data, reduction=reduction))
-        losses = torch.hstack(losses).sum(dim=-1, keepdim=True)
                 losses.backward()
                 optimizer.step()
+        losses = torch.hstack(losses).sum(dim=-1, keepdim=True)
         return losses.detach()
 
     @property