def one_shot(gpt, task):
t = gpt.training
gpt.eval()
-
+ mode='head'
+ dim_in=args.dim_model * (args.nb_blocks * 2 if mode=='deep' else 1)
model = nn.Sequential(
- nn.Linear(args.dim_model, args.dim_model),
+ nn.Linear(dim_in, args.dim_model),
nn.ReLU(),
nn.Linear(args.dim_model, args.dim_model),
nn.ReLU(),
acc_train_loss, nb_train_samples = 0, 0
for input, targets in task.policy_batches(split="train"):
- output_gpt = gpt(mygpt.BracketedSequence(input), with_readout=False).x
+ output_gpt = gpt(mygpt.BracketedSequence(input), mode=mode).x
output = model(output_gpt)
targets = targets * (input.unsqueeze(-1) == maze.v_empty)
output = output * (input.unsqueeze(-1) == maze.v_empty)
acc_test_loss, nb_test_samples = 0, 0
for input, targets in task.policy_batches(split="test"):
- output_gpt = gpt(mygpt.BracketedSequence(input), with_readout=False).x
+ output_gpt = gpt(mygpt.BracketedSequence(input), mode=mode).x
output = model(output_gpt)
targets = targets * (input.unsqueeze(-1) == maze.v_empty)
output = output * (input.unsqueeze(-1) == maze.v_empty)
# -------------------
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_gpt = gpt(mygpt.BracketedSequence(input), mode=mode).x
output = model(output_gpt)
# losses = (-output.log_softmax(-1) * targets + targets.xlogy(targets)).sum(-1)
# losses = losses * (input == maze.v_empty)