added support for a small maze dataset, with the purpose of debugging and iteration

tasks/mazes/README.md

@@ -8,3 +8,9 @@ To run the maze training that we used for the paper, run the following command f
 ```
 python -m tasks.mazes.train --d_model 2048 --d_input 512 --synapse_depth 4 --heads 8 --n_synch_out 64 --n_synch_action 32 --neuron_select_type first-last --iterations 75 --memory_length 25 --deep_memory --memory_hidden_dims 32 --dropout 0.1 --no-do_normalisation --positional_embedding_type none --backbone_type resnet34-2 --batch_size 64 --batch_size_test 64 --lr 1e-4 --training_iterations 1000001 --warmup_steps 10000 --use_scheduler --scheduler_type cosine --weight_decay 0.0 --log_dir logs/mazes/d=2048--i=512--h=8--ns=64-32--iters=75x25--h=32--drop=0.1--pos=none--back=34-2--seed=42 --dataset mazes-medium --save_every 2000 --track_every 5000 --seed 42 --n_test_batches 50
 ```
+
+## Small training run
+We also provide a 'mazes-small' dataset (see [here](https://drive.google.com/file/d/1cBgqhaUUtsrll8-o2VY42hPpyBcfFv86/view?usp=drivesdk)) for fast iteration and testing ideas. The following command can train a CTM locally without a GPU in 12-24 hours:
+```
+python -m tasks.mazes.train --dataset mazes-small --maze_route_length 50 --cirriculum_lookahead 5  --model ctm --d_model 1024 --d_input 256 --backbone_type resnet18-1 --synapse_depth 8 --heads 4 --n_synch_out 128 --n_synch_action 128 --neuron_select_type random-pairing --memory_length 25 --iterations 50 --training_iterations 100001 --lr 1e-4 --batch_size 64 --batch_size_test 32 --n_test_batches 50 --log_dir logs/mazes-small-tester --track_every 2000
+```

tasks/mazes/plotting.py

@@ -96,7 +96,7 @@ def make_maze_gif(inputs, predictions, targets, attention_tracking, save_locatio
     route_colours = []
     solution_maze = draw_path(np.moveaxis(inputs, 0, -1), targets)
     
-    # cv2.imwrite(f'{save_location}/ground_truth.png', solution_maze[:,:,::-1]*255)
+    n_heads = attention_tracking.shape[1]
     mosaic = [['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
               ['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
               ['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
@@ -104,9 +104,25 @@ def make_maze_gif(inputs, predictions, targets, attention_tracking, save_locatio
               ['head_0', 'head_1', 'head_2', 'head_3', 'head_4', 'head_5', 'head_6', 'head_7'],
               ['head_8', 'head_9', 'head_10', 'head_11', 'head_12', 'head_13', 'head_14', 'head_15'],
               ]
+    if n_heads == 8:
+        mosaic = [['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
+              ['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
+              ['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
+              ['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
+              ['head_0', 'head_1', 'head_2', 'head_3', 'head_4', 'head_5', 'head_6', 'head_7'],
+              ]
+    elif n_heads == 4:
+        mosaic = [['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
+              ['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
+              ['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
+              ['overlay', 'overlay', 'overlay', 'overlay', 'route', 'route', 'route', 'route'],
+              ['head_0', 'head_0', 'head_1', 'head_1', 'head_2', 'head_2', 'head_3', 'head_3'],
+              ['head_0', 'head_0', 'head_1', 'head_1', 'head_2', 'head_2', 'head_3', 'head_3'],
+              ]
+
     img_aspect = 1
     figscale = 1
-    aspect_ratio = (8 * figscale, 6 * figscale * img_aspect) # W, H
+    aspect_ratio = (len(mosaic[0]) * figscale, len(mosaic) * figscale * img_aspect) # W, H
     
     route_steps = [np.unravel_index(np.argmax((inputs == np.reshape(np.array([1, 0, 0]), (3, 1, 1))).all(0)), inputs.shape[1:])]  # Starting point
     frames = []

tasks/mazes/train.py

@@ -108,7 +108,7 @@ def parse_args():
 
     # Logging and Saving
     parser.add_argument('--log_dir', type=str, default='logs/scratch', help='Directory for logging.')
-    parser.add_argument('--dataset', type=str, default='mazes-medium', help='Dataset to use.', choices=['mazes-medium', 'mazes-large']) 
+    parser.add_argument('--dataset', type=str, default='mazes-medium', help='Dataset to use.', choices=['mazes-medium', 'mazes-large', 'mazes-small']) 
     parser.add_argument('--data_root', type=str, default='data/mazes', help='Data root.')
     
     parser.add_argument('--save_every', type=int, default=1000, help='Save checkpoints every this many iterations.')
@@ -139,7 +139,7 @@ if __name__=='__main__':
     set_seed(args.seed, False)
     if not os.path.exists(args.log_dir): os.makedirs(args.log_dir)
 
-    assert args.dataset in ['mazes-medium', 'mazes-large']
+    assert args.dataset in ['mazes-medium', 'mazes-large', 'mazes-small']