From 359cf44b609cebd0f01b9c2d2be1f76a4577a97b Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Sun, 7 Jan 2024 16:05:53 +0100 Subject: [PATCH] Update. --- mygpt.py | 54 ++++++++++++++++++++++++++++++------------------------ 1 file changed, 30 insertions(+), 24 deletions(-) diff --git a/mygpt.py b/mygpt.py index 21ae739..7aa8578 100755 --- a/mygpt.py +++ b/mygpt.py @@ -656,34 +656,40 @@ class Caterpillar(nn.Module): self.rec_K[:, :, t0:t1] = next_K.flatten(2, 3) if self.training and self.proba_flashback: - insert_flash_back( - self.rec_V, - V, - self.rec_K, - K, - t0, - t1, - CL, - proba=self.proba_flashback / CL, - ) + # insert_flash_back( + # self.rec_V, + # V, + # self.rec_K, + # K, + # t0, + # t1, + # CL, + # proba=self.proba_flashback / CL, + # ) - # n = torch.arange(N, device=X.device)[:, None, None, None] - # t = torch.arange(t0, t1, device=X.device)[None, None, :, None] - # dv = torch.arange(DV)[None, None, None, :] - # dk = torch.arange(DK)[None, None, None, :] + n = torch.arange(N, device=X.device)[:, None, None, None] + t = torch.arange(t0, t1, device=X.device)[None, None, :, None] + dv = torch.arange(DV)[None, None, None, :] + dk = torch.arange(DK)[None, None, None, :] - # u = ( - # torch.rand(N, CH, t1 - t0, 1, device=X.device).mul(t).long() // CL - # ) * CL + u = ( + torch.rand(N, CH, t1 - t0, 1, device=X.device).mul(t).long() // CL + ) * CL - # src_time = t - u - t0 - # src_head = torch.randint(H, (N, CH, t1 - t0, 1), device=X.device) + src_time = t - u - t0 + src_head = torch.randint(H, (N, CH, t1 - t0, 1), device=X.device) - # mk = ( - # torch.rand(self.rec_V[:, :, t0:t1].size()) <= self.proba_flashback - # ).long() - # self.rec_V[:, :, t0:t1] = V[n, src_head, src_time, dv] - # self.rec_K[:, :, t0:t1] = K[n, src_head, src_time, dk] + mask_V = (torch.rand(N, CH, t1 - t0, DV) <= self.proba_flashback).long() + self.rec_V[:, :, t0:t1] = ( + mask_V * V[n, src_head, src_time, dv] + + (1 - mask_V) * self.rec_V[:, :, t0:t1] + ) + + mask_K = (torch.rand(N, CH, t1 - t0, DK) <= self.proba_flashback).long() + self.rec_K[:, :, t0:t1] = ( + mask_K * K[n, src_head, src_time, dk] + + (1 - mask_K) * self.rec_K[:, :, t0:t1] + ) exit(0) -- 2.39.5