From 41c7509dc3d2153da79ed09ecf4a3b592503f15e Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Fri, 17 Mar 2023 17:22:40 +0100 Subject: [PATCH] Update --- beaver.py | 28 ++++++++++++++++++++++++---- maze.py | 37 ++++++++++++++++++++++++++++--------- 2 files changed, 52 insertions(+), 13 deletions(-) diff --git a/beaver.py b/beaver.py index 54510f0..2cc2140 100755 --- a/beaver.py +++ b/beaver.py @@ -178,7 +178,10 @@ def one_shot(gpt, task): nn.Linear(args.dim_model, 4), ).to(device) + print(f"{args.nb_epochs=}") + for n_epoch in range(args.nb_epochs): + print(f"{n_epoch=}") learning_rate = learning_rate_schedule[n_epoch] optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) @@ -212,6 +215,23 @@ def one_shot(gpt, task): f"diff_ce {n_epoch} train {acc_train_loss/nb_train_samples} test {acc_test_loss/nb_test_samples}" ) + # ------------------- + input, targets = next(task.policy_batches(split="test")) + output_gpt = gpt(mygpt.BracketedSequence(input), with_readout=False).x + output = model(output_gpt) + losses = (-output.log_softmax(-1) * targets + targets.xlogy(targets)).sum(-1) + losses = losses / losses.max() + print(f"{input.size()=} {losses.size()=} {losses.min()=} {losses.max()=}") + losses = losses * (input == 0) + losses = losses.reshape(-1, args.maze_height, args.maze_width) + input = input.reshape(-1, args.maze_height, args.maze_width) + maze.save_image( + os.path.join(args.result_dir, f"oneshot_{n_epoch:04d}.png"), + mazes=input, + score_paths=losses, + ) + # ------------------- + gpt.train(t) @@ -354,10 +374,10 @@ class TaskMaze(Task): _, predicted_paths = self.seq2map(result) maze.save_image( os.path.join(args.result_dir, f"result_{n_epoch:04d}.png"), - mazes, - paths, - predicted_paths, - maze.path_correctness(mazes, predicted_paths), + mazes=mazes, + target_paths=paths, + predicted_paths=predicted_paths, + path_correct=maze.path_correctness(mazes, predicted_paths), ) model.train(t) diff --git a/maze.py b/maze.py index 6c3fe94..6e8e179 100755 --- a/maze.py +++ b/maze.py @@ -187,9 +187,14 @@ def create_maze_data( ###################################################################### -def save_image(name, mazes, target_paths, predicted_paths=None, path_correct=None): - mazes, target_paths = mazes.cpu(), target_paths.cpu() - +def save_image( + name, + mazes, + target_paths=None, + predicted_paths=None, + score_paths=None, + path_correct=None, +): colors = torch.tensor( [ [255, 255, 255], # empty @@ -200,13 +205,22 @@ def save_image(name, mazes, target_paths, predicted_paths=None, path_correct=Non ] ) + mazes = mazes.cpu() + mazes = colors[mazes.reshape(-1)].reshape(mazes.size() + (-1,)).permute(0, 3, 1, 2) - target_paths = ( - colors[target_paths.reshape(-1)] - .reshape(target_paths.size() + (-1,)) - .permute(0, 3, 1, 2) - ) - imgs = torch.cat((mazes.unsqueeze(1), target_paths.unsqueeze(1)), 1) + + imgs = mazes.unsqueeze(1) + + if target_paths is not None: + target_paths = target_paths.cpu() + + target_paths = ( + colors[target_paths.reshape(-1)] + .reshape(target_paths.size() + (-1,)) + .permute(0, 3, 1, 2) + ) + + imgs = torch.cat((imgs, target_paths.unsqueeze(1)), 1) if predicted_paths is not None: predicted_paths = predicted_paths.cpu() @@ -217,6 +231,11 @@ def save_image(name, mazes, target_paths, predicted_paths=None, path_correct=Non ) imgs = torch.cat((imgs, predicted_paths.unsqueeze(1)), 1) + if score_paths is not None: + score_paths = (score_paths.cpu() * 255.0).long() + score_paths = score_paths.unsqueeze(1).expand(-1, 3, -1, -1) + imgs = torch.cat((imgs, score_paths.unsqueeze(1)), 1) + # NxKxCxHxW if path_correct is None: path_correct = torch.zeros(imgs.size(0)) <= 1 -- 2.39.5