X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=beaver.py;h=f850f699ed9c3d18ef521df99b0a77595485fa7c;hb=e14c55948c1f099f95fa3d7343b5c939e60fcb1c;hp=4f41832fcd928b7f25bc9722a66b2d810229711e;hpb=eb86c22964f03f186ee225f129bc260128b10f9a;p=beaver.git

diff --git a/beaver.py b/beaver.py
index 4f41832..f850f69 100755
--- a/beaver.py
+++ b/beaver.py
@@ -205,7 +205,11 @@ def compute_perplexity(model, task, fixed_len, split="train"):
         for input in task.batches(split=split):
             input = input.to(device)
             output = eval_mygpt(model, input, fixed_len=fixed_len)
-            loss = F.cross_entropy(output.transpose(1, 2), input)
+            if args.noncausal_prompt:
+                t = input.size(1) // 2
+                loss = F.cross_entropy(output[:, t:].transpose(1, 2), input[:, t:])
+            else:
+                loss = F.cross_entropy(output.transpose(1, 2), input)
             acc_loss += loss.item() * input.size(0)
             nb_samples += input.size(0)
 
@@ -524,7 +528,10 @@ amm_generator = None
 if args.noncausal_prompt:
     amm_generator = lambda d: torch.logical_and(
         torch.arange(d)[None, None, :, None] < torch.arange(d)[None, None, None, :],
-        torch.arange(d)[None, None, :, None] >= d // 2,
+        torch.logical_or(
+            torch.arange(d)[None, None, :, None] >= d // 2,
+            torch.arange(d)[None, None, None, :] >= d // 2,
+        ),
     )
 
 model = mygpt.MyGPT(
@@ -645,7 +652,11 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs):
         output = eval_mygpt(
             model, input, mode=args.oneshot_input, fixed_len=task.height * task.width
         )
-        loss = F.cross_entropy(output.transpose(1, 2), input)
+        if args.noncausal_prompt:
+            t = input.size(1) // 2
+            loss = F.cross_entropy(output[:, t:].transpose(1, 2), input[:, t:])
+        else:
+            loss = F.cross_entropy(output.transpose(1, 2), input)
         acc_train_loss += loss.item() * input.size(0)
         nb_train_samples += input.size(0)