From a0e547917131af0b353e3bf31a062c9b35c8dd18 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Fri, 17 Mar 2023 22:15:40 +0100 Subject: [PATCH] Update --- beaver.py | 25 +++++++++++++++---------- maze.py | 4 ++-- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/beaver.py b/beaver.py index 2cc2140..c3b7e09 100755 --- a/beaver.py +++ b/beaver.py @@ -173,15 +173,14 @@ def one_shot(gpt, task): t = gpt.training gpt.eval() model = nn.Sequential( + nn.Linear(args.dim_model, args.dim_model), + nn.ReLU(), nn.Linear(args.dim_model, args.dim_model), nn.ReLU(), 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) @@ -189,9 +188,12 @@ def one_shot(gpt, task): for input, targets in task.policy_batches(split="train"): output_gpt = gpt(mygpt.BracketedSequence(input), with_readout=False).x output = model(output_gpt) + targets = targets * (input.unsqueeze(-1) == maze.v_empty) + output = output * (input.unsqueeze(-1) == maze.v_empty) loss = ( - -(output.log_softmax(-1) * targets).sum(-1).mean() - + targets.xlogy(targets).sum(-1).mean() + -(output.log_softmax(-1) * targets).sum() + / (input == maze.v_empty).sum() + + targets.xlogy(targets).sum() / (input == maze.v_empty).sum() ) acc_train_loss += loss.item() * input.size(0) nb_train_samples += input.size(0) @@ -204,9 +206,12 @@ def one_shot(gpt, task): for input, targets in task.policy_batches(split="test"): output_gpt = gpt(mygpt.BracketedSequence(input), with_readout=False).x output = model(output_gpt) + targets = targets * (input.unsqueeze(-1) == maze.v_empty) + output = output * (input.unsqueeze(-1) == maze.v_empty) loss = ( - -(output.log_softmax(-1) * targets).sum(-1).mean() - + targets.xlogy(targets).sum(-1).mean() + -(output.log_softmax(-1) * targets).sum() + / (input == maze.v_empty).sum() + + targets.xlogy(targets).sum() / (input == maze.v_empty).sum() ) acc_test_loss += loss.item() * input.size(0) nb_test_samples += input.size(0) @@ -216,13 +221,13 @@ def one_shot(gpt, task): ) # ------------------- - input, targets = next(task.policy_batches(split="test")) + input = task.test_input[:32, : task.height * task.width] + targets = task.test_policies[:32] 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 * (input == maze.v_empty) 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( diff --git a/maze.py b/maze.py index 754cdea..44bef7c 100755 --- a/maze.py +++ b/maze.py @@ -200,7 +200,7 @@ def save_image( [255, 255, 255], # empty [0, 0, 0], # wall [0, 255, 0], # start - [0, 0, 255], # goal + [127, 127, 255], # goal [255, 0, 0], # path ] ) @@ -238,7 +238,7 @@ def save_image( c_score_paths = score_paths.unsqueeze(1).expand(-1, 3, -1, -1) c_score_paths = ( c_score_paths * colors[4].reshape(1, 3, 1, 1) - + (1 - c_score_paths) * colors[3].reshape(1, 3, 1, 1) + + (1 - c_score_paths) * colors[0].reshape(1, 3, 1, 1) ).long() c_score_paths = c_score_paths * (mazes.unsqueeze(1) == v_empty) + c_mazes * ( mazes.unsqueeze(1) != v_empty -- 2.20.1