From 910434daa74525b29d7ba117312fc67789d1ad84 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Fri, 14 Jul 2023 18:08:35 +0200 Subject: [PATCH] Update. --- world.py | 63 +++++++++++++++++++++++++++++++++++--------------------- 1 file changed, 39 insertions(+), 24 deletions(-) diff --git a/world.py b/world.py index b077987..61a07e9 100755 --- a/world.py +++ b/world.py @@ -146,7 +146,7 @@ def train_encoder( for input in tqdm.tqdm(train_input.split(batch_size), desc="vqae-train"): z = encoder(input) - zq = z if k < 1 else quantizer(z) + zq = z if k < 2 else quantizer(z) output = decoder(zq) output = output.reshape( @@ -334,8 +334,6 @@ def create_data_and_processors(nb_train_samples, nb_test_samples, nb_epochs=10): train_input = generate_episodes(nb_train_samples, steps) test_input = generate_episodes(nb_test_samples, steps) - print(f"{train_input.size()=} {test_input.size()=}") - encoder, quantizer, decoder = train_encoder( train_input, test_input, nb_epochs=nb_epochs ) @@ -347,26 +345,40 @@ def create_data_and_processors(nb_train_samples, nb_test_samples, nb_epochs=10): pow2 = (2 ** torch.arange(z.size(1), device=z.device))[None, None, :] z_h, z_w = z.size(2), z.size(3) - def frame2seq(x): - z = encoder(x) - ze_bool = (quantizer(z) >= 0).long() - seq = ( - ze_bool.permute(0, 2, 3, 1).reshape(ze_bool.size(0), -1, ze_bool.size(1)) - * pow2 - ).sum(-1) - return seq - - def seq2frame(seq, T=1e-2): - zd_bool = (seq[:, :, None] // pow2) % 2 - zd_bool = zd_bool.reshape(zd_bool.size(0), z_h, z_w, -1).permute(0, 3, 1, 2) - logits = decoder(zd_bool * 2.0 - 1.0) - logits = logits.reshape( - logits.size(0), -1, 3, logits.size(2), logits.size(3) - ).permute(0, 2, 3, 4, 1) - results = torch.distributions.categorical.Categorical( - logits=logits / T - ).sample() - return results + def frame2seq(input, batch_size=25): + seq = [] + + for x in input.split(batch_size): + z = encoder(x) + ze_bool = (quantizer(z) >= 0).long() + output = ( + ze_bool.permute(0, 2, 3, 1).reshape( + ze_bool.size(0), -1, ze_bool.size(1) + ) + * pow2 + ).sum(-1) + + seq.append(output) + + return torch.cat(seq, dim=0) + + def seq2frame(input, batch_size=25, T=1e-2): + frames = [] + + for seq in input.split(batch_size): + zd_bool = (seq[:, :, None] // pow2) % 2 + zd_bool = zd_bool.reshape(zd_bool.size(0), z_h, z_w, -1).permute(0, 3, 1, 2) + logits = decoder(zd_bool * 2.0 - 1.0) + logits = logits.reshape( + logits.size(0), -1, 3, logits.size(2), logits.size(3) + ).permute(0, 2, 3, 4, 1) + output = torch.distributions.categorical.Categorical( + logits=logits / T + ).sample() + + frames.append(output) + + return torch.cat(frames, dim=0) return train_input, test_input, frame2seq, seq2frame @@ -375,7 +387,10 @@ def create_data_and_processors(nb_train_samples, nb_test_samples, nb_epochs=10): if __name__ == "__main__": train_input, test_input, frame2seq, seq2frame = create_data_and_processors( - 10000, 1000 + # 10000, 1000, + 100, + 100, + nb_epochs=2, ) input = test_input[:64] -- 2.39.5