Update

[beaver.git] / beaver.py

diff --git a/beaver.py b/beaver.py
index bd17365..86008f6 100755 (executable)
--- a/beaver.py
+++ b/beaver.py
@@ -134,7 +134,7 @@ for n in vars(args):
 def generation_order(x, fixed_len):
     if args.random_regression_order:
         order = torch.rand(x.size(), device=x.device)
-        order[:, :fixed_len] = torch.linspace(-2, -1, fixed_len, device=order.device)
+        order[:, :fixed_len] = torch.linspace(-2, -1, fixed_len, device=x.device)
         order = order.sort(1).indices
     else:
         order = (
@@ -143,10 +143,10 @@ def generation_order(x, fixed_len):
     return order
 
 
-def reorder(x, order, back=False):  # x is NxTxD1x...xDk, order is NxT'
+def reorder(x, order, reverse=False):  # x is NxTxD1x...xDk, order is NxT'
     u = x.reshape(x.size()[:2] + (-1,))
     order = order.unsqueeze(-1).expand(-1, -1, u.size(-1))
-    if back:
+    if reverse:
         v = u.new(u.size())
         v.scatter_(1, order, u)
     else:
@@ -167,7 +167,9 @@ def shuffle(x, fixed_len):
 
 
 def masked_inplace_autoregression(model, batch_size, input, ar_mask, order=None):
-    for input, ar_mask in zip(input.split(batch_size), ar_mask.split(batch_size)):
+    for input, ar_mask, order in zip(
+        input.split(batch_size), ar_mask.split(batch_size), order.split(batch_size)
+    ):
         i = (ar_mask.sum(0) > 0).nonzero()
         if i.min() > 0:
             # Needed to initialize the model's cache
@@ -186,7 +188,7 @@ def masked_inplace_autoregression(model, batch_size, input, ar_mask, order=None)
 ######################################################################
 
 
-def compute_perplexity(model, fixed_len, split="train"):
+def compute_perplexity(model, task, fixed_len, split="train"):
     with torch.autograd.no_grad():
         t = model.training
         model.eval()
@@ -197,7 +199,7 @@ def compute_perplexity(model, fixed_len, split="train"):
             input = input.to(device)
             x, order = shuffle(input, fixed_len)
             x = model(mygpt.BracketedSequence(x), order=order).x
-            output = reorder(x, order, back=True)
+            output = reorder(x, order, reverse=True)
             loss = F.cross_entropy(output.transpose(1, 2), input)
             acc_loss += loss.item() * input.size(0)
             nb_samples += input.size(0)
@@ -263,7 +265,7 @@ def oneshot(gpt, task):
         for mazes, policies in task.policy_batches(split="train"):
             x, order = shuffle(mazes, task.height * task.width)
             x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x
-            output_gpt = reorder(x, order, back=True)
+            output_gpt = reorder(x, order, reverse=True)
             output = model(output_gpt)
 
             loss = compute_loss(mazes, output, policies, task.height, task.width)
@@ -278,7 +280,7 @@ def oneshot(gpt, task):
         for mazes, policies in task.policy_batches(split="test"):
             x, order = shuffle(mazes, task.height * task.width)
             x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x
-            output_gpt = reorder(x, order, back=True)
+            output_gpt = reorder(x, order, reverse=True)
             output = model(output_gpt)
             loss = compute_loss(mazes, output, policies, task.height, task.width)
             acc_test_loss += loss.item() * mazes.size(0)
@@ -293,7 +295,7 @@ def oneshot(gpt, task):
         policies = task.test_policies[:32]
         x, order = shuffle(mazes, task.height * task.width)
         x = gpt(mygpt.BracketedSequence(x), mode=args.oneshot_input, order=order).x
-        output_gpt = reorder(x, order, back=True)
+        output_gpt = reorder(x, order, reverse=True)
         output = model(output_gpt)
         if args.oneshot_output == "policy":
             targets = policies.permute(0, 2, 1)
@@ -438,7 +440,7 @@ class TaskMaze(Task):
             masked_inplace_autoregression(
                 model, self.batch_size, x, ar_mask, order=order
             )
-            result = reorder(x, order, back=True)
+            result = reorder(x, order, reverse=True)
             mazes, paths = self.seq2map(result)
             nb_correct += maze.path_correctness(mazes, paths).long().sum()
             nb_total += mazes.size(0)
@@ -469,7 +471,11 @@ class TaskMaze(Task):
             ar_mask = result.new_zeros(result.size())
             ar_mask[:, self.height * self.width :] = 1
             result *= 1 - ar_mask
-            masked_inplace_autoregression(model, self.batch_size, result, ar_mask)
+            x, order = shuffle(result, self.height * self.width)
+            masked_inplace_autoregression(
+                model, self.batch_size, x, ar_mask, order=order
+            )
+            result = reorder(x, order, reverse=True)
 
             mazes, paths = self.seq2map(input)
             _, predicted_paths = self.seq2map(result)
@@ -588,10 +594,10 @@ log_string(f"learning_rate_schedule {learning_rate_schedule}")
 if nb_epochs_finished >= args.nb_epochs:
     n_epoch = nb_epochs_finished
     train_perplexity = compute_perplexity(
-        model, fixed_len=task.height * task.width, split="train"
+        model, task, fixed_len=task.height * task.width, split="train"
     )
     test_perplexity = compute_perplexity(
-        model, fixed_len=task.height * task.width, split="test"
+        model, task, fixed_len=task.height * task.width, split="test"
     )
 
     log_string(
@@ -624,7 +630,7 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs):
         input = input.to(device)
         x, order = shuffle(input, task.height * task.width)
         x = model(mygpt.BracketedSequence(x), order=order).x
-        output = reorder(x, order, back=True)
+        output = reorder(x, order, reverse=True)
         loss = F.cross_entropy(output.transpose(1, 2), input)
         acc_train_loss += loss.item() * input.size(0)
         nb_train_samples += input.size(0)
@@ -635,7 +641,7 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs):
 
     train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples))
     test_perplexity = compute_perplexity(
-        model, fixed_len=task.height * task.width, split="test"
+        model, task, fixed_len=task.height * task.width, split="test"
     )
 
     log_string(