efcc0ddd20c248335aff7acbc2f1d5caa04d607c
[picoclvr.git] / main.py
1 #!/usr/bin/env python
2
3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
5
6 # Written by Francois Fleuret <francois@fleuret.org>
7
8 # torch.backends.cuda.matmul.allow_tf23
9 # torch.autocast(torch.bfloat16)
10
11 import math, sys, argparse, time, tqdm, os
12
13 import torch, torchvision
14 from torch import nn
15 from torch.nn import functional as F
16
17 import ffutils
18 import mygpt, tasks
19
20 ######################################################################
21
22 if torch.cuda.is_available():
23     device = torch.device("cuda")
24     torch.backends.cuda.matmul.allow_tf32 = True
25 else:
26     device = torch.device("cpu")
27
28 ######################################################################
29
30 parser = argparse.ArgumentParser(
31     description="An implementation of GPT with cache.",
32     formatter_class=argparse.ArgumentDefaultsHelpFormatter,
33 )
34
35 parser.add_argument(
36     "--task",
37     type=str,
38     default="sandbox",
39     help="sandbox, picoclvr, mnist, maze, snake, stack, expr, world",
40 )
41
42 parser.add_argument("--log_filename", type=str, default="train.log", help=" ")
43
44 parser.add_argument("--result_dir", type=str, default=None)
45
46 parser.add_argument("--seed", type=int, default=0)
47
48 parser.add_argument("--nb_epochs", type=int, default=None)
49
50 parser.add_argument("--batch_size", type=int, default=None)
51
52 parser.add_argument("--nb_train_samples", type=int, default=None)
53
54 parser.add_argument("--nb_test_samples", type=int, default=None)
55
56 parser.add_argument("--optim", type=str, default="adam")
57
58 parser.add_argument("--learning_rate", type=float, default=1e-4)
59
60 parser.add_argument("--learning_rate_schedule", type=str, default="10: 2e-5,30: 4e-6")
61
62 parser.add_argument("--model", type=str, default="37M")
63
64 parser.add_argument("--dim_model", type=int, default=None)
65
66 parser.add_argument("--dim_keys", type=int, default=None)
67
68 parser.add_argument("--dim_hidden", type=int, default=None)
69
70 parser.add_argument("--nb_heads", type=int, default=None)
71
72 parser.add_argument("--nb_blocks", type=int, default=None)
73
74 parser.add_argument("--dropout", type=float, default=0.1)
75
76 parser.add_argument("--deterministic_synthesis", action="store_true", default=False)
77
78 parser.add_argument("--no_checkpoint", action="store_true", default=False)
79
80 parser.add_argument("--overwrite_results", action="store_true", default=False)
81
82 parser.add_argument("--checkpoint_name", type=str, default="checkpoint.pth")
83
84 ##############################
85 # picoclvr options
86
87 parser.add_argument("--sandbox_level", type=int, default=0)
88
89 parser.add_argument("--sandbox_levels_nb_items", type=int, default=25)
90
91 parser.add_argument("--sandbox_levels_len_source", type=int, default=6)
92
93 parser.add_argument("--sandbox_levels_len_result", type=int, default=8)
94
95 ##############################
96 # picoclvr options
97
98 parser.add_argument("--picoclvr_nb_colors", type=int, default=5)
99
100 parser.add_argument("--picoclvr_height", type=int, default=12)
101
102 parser.add_argument("--picoclvr_width", type=int, default=16)
103
104 parser.add_argument("--picocvlr_prune_properties", type=str, default="none")
105
106 ##############################
107 # Maze options
108
109 parser.add_argument("--maze_height", type=int, default=23)
110
111 parser.add_argument("--maze_width", type=int, default=39)
112
113 parser.add_argument("--maze_nb_walls", type=int, default=45)
114
115 ##############################
116 # Snake options
117
118 parser.add_argument("--snake_height", type=int, default=6)
119
120 parser.add_argument("--snake_width", type=int, default=8)
121
122 parser.add_argument("--snake_nb_colors", type=int, default=5)
123
124 parser.add_argument("--snake_length", type=int, default=200)
125
126 ##############################
127 # Stack options
128
129 parser.add_argument("--stack_nb_steps", type=int, default=100)
130
131 parser.add_argument("--stack_nb_stacks", type=int, default=3)
132
133 parser.add_argument("--stack_nb_digits", type=int, default=3)
134
135 parser.add_argument("--stack_fraction_values_for_train", type=float, default=0.75)
136
137 ##############################
138 # Expr options
139
140 parser.add_argument("--expr_nb_variables", type=int, default=5)
141
142 parser.add_argument("--expr_sequence_length", type=int, default=40)
143
144 parser.add_argument("--expr_operand_max", type=int, default=9)
145
146 parser.add_argument("--expr_result_max", type=int, default=99)
147
148 parser.add_argument("--expr_input_file", type=str, default=None)
149
150 ##############################
151 # World options
152
153 parser.add_argument("--world_vqae_nb_epochs", type=int, default=25)
154
155 ######################################################################
156
157 args = parser.parse_args()
158
159 assert args.picocvlr_prune_properties in {"none", "train+eval", "eval"}
160
161 if args.result_dir is None:
162     args.result_dir = f"results_{args.task}"
163
164 ######################################################################
165
166 default_task_args = {
167     "sandbox": {
168         "nb_epochs": 50,
169         "batch_size": 25,
170         "nb_train_samples": 100000,
171         "nb_test_samples": 10000,
172     },
173     "picoclvr": {
174         "nb_epochs": 25,
175         "batch_size": 25,
176         "nb_train_samples": 250000,
177         "nb_test_samples": 10000,
178     },
179     "mnist": {
180         "nb_epochs": 25,
181         "batch_size": 10,
182         "nb_train_samples": 250000,
183         "nb_test_samples": 10000,
184     },
185     "maze": {
186         "nb_epochs": 25,
187         "batch_size": 5,
188         "nb_train_samples": 250000,
189         "nb_test_samples": 10000,
190     },
191     "snake": {
192         "nb_epochs": 5,
193         "batch_size": 25,
194         "nb_train_samples": 250000,
195         "nb_test_samples": 10000,
196     },
197     "stack": {
198         "nb_epochs": 5,
199         "batch_size": 25,
200         "nb_train_samples": 100000,
201         "nb_test_samples": 1000,
202     },
203     "expr": {
204         "nb_epochs": 40,
205         "batch_size": 25,
206         "nb_train_samples": 1000000,
207         "nb_test_samples": 10000,
208     },
209     "world": {
210         "nb_epochs": 10,
211         "batch_size": 25,
212         "nb_train_samples": 25000,
213         "nb_test_samples": 1000,
214     },
215 }
216
217 if args.task in default_task_args:
218     for k, v in default_task_args[args.task].items():
219         if getattr(args, k) is None:
220             setattr(args, k, v)
221
222 ######################################################################
223
224 default_model_args = {
225     "17K": {
226         "dim_model": 32,
227         "dim_keys": 32,
228         "dim_hidden": 32,
229         "nb_heads": 2,
230         "nb_blocks": 2,
231     },
232     "37M": {
233         "dim_model": 512,
234         "dim_keys": 64,
235         "dim_hidden": 2048,
236         "nb_heads": 8,
237         "nb_blocks": 12,
238     },
239     "122M": {
240         "dim_model": 768,
241         "dim_keys": 64,
242         "dim_hidden": 2048,
243         "nb_heads": 8,
244         "nb_blocks": 24,
245     },
246     "352M": {
247         "dim_model": 1024,
248         "dim_keys": 64,
249         "dim_hidden": 2048,
250         "nb_heads": 8,
251         "nb_blocks": 48,
252     },
253 }
254
255 if args.model in default_model_args:
256     for k, v in default_model_args[args.model].items():
257         if getattr(args, k) is None:
258             setattr(args, k, v)
259 else:
260     raise ValueError(f"Unknown model {args.model}")
261
262 ######################################################################
263
264 try:
265     os.mkdir(args.result_dir)
266 except FileExistsError:
267     if not args.overwrite_results:
268         print(f"result directory {args.result_dir} already exists")
269         exit(1)
270
271 log_file = open(os.path.join(args.result_dir, args.log_filename), "a")
272
273 if args.seed >= 0:
274     # torch.backends.cudnn.deterministic = True
275     # torch.backends.cudnn.benchmark = False
276     # torch.use_deterministic_algorithms(True)
277     torch.manual_seed(args.seed)
278     if torch.cuda.is_available():
279         torch.cuda.manual_seed_all(args.seed)
280
281 ######################################################################
282
283
284 def log_string(s):
285     t = time.strftime("%Y%m%d-%H:%M:%S ", time.localtime())
286
287     if log_file is not None:
288         log_file.write(t + s + "\n")
289         log_file.flush()
290
291     print(t + s)
292     sys.stdout.flush()
293
294
295 for n in vars(args):
296     log_string(f"args.{n} {getattr(args, n)}")
297
298
299 ######################################################################
300
301
302 def picoclvr_pruner_horizontal_green(p):
303     return not ("green" in p and ("left" in p or "right" in p))
304
305
306 picoclvr_pruner_train = (
307     picoclvr_pruner_horizontal_green
308     if args.picocvlr_prune_properties in {"train+eval"}
309     else None
310 )
311
312 picoclvr_pruner_eval = (
313     (lambda p: not picoclvr_pruner_horizontal_green(p))
314     if args.picocvlr_prune_properties in {"train+eval", "eval"}
315     else None
316 )
317
318 ######################################################################
319
320 if args.task == "sandbox":
321     if args.sandbox_level == 0:
322         problem = tasks.ProblemLevel0(
323             nb_sentences=args.sandbox_levels_nb_items,
324             len_prompt=args.sandbox_levels_len_source,
325             len_result=args.sandbox_levels_len_result,
326         )
327     elif args.sandbox_level == 1:
328         problem = tasks.ProblemLevel1(
329             nb_operators=args.sandbox_levels_nb_items,
330             len_source=args.sandbox_levels_len_source,
331             len_result=args.sandbox_levels_len_result,
332         )
333     elif args.sandbox_level == 2:
334         problem = tasks.ProblemLevel2(
335             len_source=args.sandbox_levels_len_source,
336             len_result=args.sandbox_levels_len_result,
337         )
338     else:
339         raise ValueError(f"Unknown sandbox level {args.sandbox_level}")
340
341     task = tasks.SandBox(
342         problem,
343         # tasks.ProblemAddition(zero_padded=False, inverted_result=False),
344         nb_train_samples=args.nb_train_samples,
345         nb_test_samples=args.nb_test_samples,
346         batch_size=args.batch_size,
347         logger=log_string,
348         device=device,
349     )
350
351 elif args.task == "picoclvr":
352     task = tasks.PicoCLVR(
353         nb_train_samples=args.nb_train_samples,
354         nb_test_samples=args.nb_test_samples,
355         batch_size=args.batch_size,
356         height=args.picoclvr_height,
357         width=args.picoclvr_width,
358         nb_colors=args.picoclvr_nb_colors,
359         logger=log_string,
360         device=device,
361         pruner_train=picoclvr_pruner_train,
362         pruner_eval=picoclvr_pruner_eval,
363     )
364
365 elif args.task == "mnist":
366     task = tasks.MNIST(
367         nb_train_samples=args.nb_train_samples,
368         nb_test_samples=args.nb_test_samples,
369         batch_size=args.batch_size,
370         device=device,
371     )
372
373 elif args.task == "maze":
374     task = tasks.Maze(
375         nb_train_samples=args.nb_train_samples,
376         nb_test_samples=args.nb_test_samples,
377         batch_size=args.batch_size,
378         height=args.maze_height,
379         width=args.maze_width,
380         nb_walls=args.maze_nb_walls,
381         device=device,
382     )
383
384 elif args.task == "snake":
385     task = tasks.Snake(
386         nb_train_samples=args.nb_train_samples,
387         nb_test_samples=args.nb_test_samples,
388         batch_size=args.batch_size,
389         height=args.snake_height,
390         width=args.snake_width,
391         nb_colors=args.snake_nb_colors,
392         length=args.snake_length,
393         prompt_length=args.snake_length // 2,
394         device=device,
395     )
396
397 elif args.task == "stack":
398     task = tasks.Stack(
399         nb_train_samples=args.nb_train_samples,
400         nb_test_samples=args.nb_test_samples,
401         batch_size=args.batch_size,
402         logger=log_string,
403         nb_steps=args.stack_nb_steps,
404         nb_stacks=args.stack_nb_stacks,
405         nb_digits=args.stack_nb_digits,
406         fraction_values_for_train=args.stack_fraction_values_for_train,
407         device=device,
408     )
409
410 elif args.task == "expr":
411     task = tasks.Expr(
412         nb_train_samples=args.nb_train_samples,
413         nb_test_samples=args.nb_test_samples,
414         nb_variables=args.expr_nb_variables,
415         sequence_length=args.expr_sequence_length,
416         operand_max=args.expr_operand_max,
417         result_max=args.expr_result_max,
418         batch_size=args.batch_size,
419         device=device,
420     )
421
422 elif args.task == "world":
423     task = tasks.World(
424         nb_train_samples=args.nb_train_samples,
425         nb_test_samples=args.nb_test_samples,
426         batch_size=args.batch_size,
427         vqae_nb_epochs=args.world_vqae_nb_epochs,
428         logger=log_string,
429         device=device,
430     )
431
432 else:
433     raise ValueError(f"Unknown task {args.task}")
434
435 ######################################################################
436
437 log_string(f"device {device}")
438
439 vocabulary_size = task.vocabulary_size()
440
441 log_string(f"vocabulary_size {vocabulary_size}")
442
443 ##############################
444
445 model = mygpt.MyGPT(
446     vocabulary_size=vocabulary_size,
447     dim_model=args.dim_model,
448     dim_keys=args.dim_keys,
449     dim_hidden=args.dim_hidden,
450     nb_heads=args.nb_heads,
451     nb_blocks=args.nb_blocks,
452     causal=True,
453     dropout=args.dropout,
454 )
455
456 model.to(device)
457
458 nb_parameters = sum(p.numel() for p in model.parameters())
459 log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)")
460
461 ######################################################################
462
463 nb_epochs_finished = 0
464
465 if args.no_checkpoint:
466     log_string(f"not trying to load checkpoint.")
467
468 else:
469     try:
470         checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name)
471         checkpoint = torch.load(checkpoint_name)
472         nb_epochs_finished = checkpoint["nb_epochs_finished"]
473         model.load_state_dict(checkpoint["model_state"])
474         torch.set_rng_state(checkpoint["rng_state"])
475         if torch.cuda.is_available():
476             torch.cuda.set_rng_state(checkpoint["cuda_rng_state"])
477
478         log_string(f"checkpoint loaded with {nb_epochs_finished} epochs finished.")
479
480     except FileNotFoundError:
481         log_string("starting from scratch.")
482
483     except:
484         log_string("error when loading the checkpoint.")
485         exit(1)
486
487 ######################################################################
488
489 if args.task == "expr" and args.expr_input_file is not None:
490     task.produce_results(
491         nb_epochs_finished,
492         model,
493         args.result_dir,
494         log_string,
495         args.deterministic_synthesis,
496         args.expr_input_file,
497     )
498
499     exit(0)
500
501 ######################################################################
502
503 nb_epochs = args.nb_epochs if args.nb_epochs > 0 else nb_epochs_default
504
505 # Compute the entropy of the training tokens
506
507 token_count = 0
508 for input in task.batches(split="train"):
509     token_count += F.one_hot(input, num_classes=task.vocabulary_size()).sum((0, 1))
510 token_probas = token_count / token_count.sum()
511 entropy = -torch.xlogy(token_probas, token_probas).sum()
512 train_set_perplexity = math.exp(entropy)
513
514 ##############################
515
516 # A bit of paranoia never hurts
517
518 train_examples = {}
519
520
521 for input in task.batches(split="train"):
522     assert input.dim() == 2 and input.dtype == torch.int64
523     for x in input:
524         train_examples[x.sum().item()] = x
525
526 nb_total, nb_collisions = 0, 0
527 for input in task.batches(split="test"):
528     assert input.dim() == 2 and input.dtype == torch.int64
529     for x in input:
530         nb_total += 1
531         y = train_examples.get(x.sum().item())
532         if y is not None:
533             if x.size() == y.size() and (x - y).abs().sum() == 0:
534                 nb_collisions += 1
535
536 del train_examples
537
538 log_string(
539     f"data_check {nb_collisions*100/nb_total:.02f}% ({nb_collisions}/{nb_total}) of test samples are in the train set"
540 )
541
542 ##############################
543
544 if args.learning_rate_schedule == "cos":
545     learning_rate_schedule = {}
546     for n_epoch in range(args.nb_epochs):
547         u = n_epoch / args.nb_epochs * math.pi
548         learning_rate_schedule[n_epoch] = args.learning_rate * 0.5 * (1 + math.cos(u))
549 else:
550     u = {
551         int(k): float(v)
552         for k, v in [
553             tuple(x.split(":")) for x in args.learning_rate_schedule.split(",")
554         ]
555     }
556
557     learning_rate_schedule = {}
558     learning_rate = args.learning_rate
559     for n_epoch in range(args.nb_epochs):
560         if n_epoch in u:
561             learning_rate = u[n_epoch]
562         learning_rate_schedule[n_epoch] = learning_rate
563
564 log_string(f"learning_rate_schedule {learning_rate_schedule}")
565
566 ##############################
567
568 nb_samples_seen = 0
569
570 if nb_epochs_finished >= nb_epochs:
571     task.produce_results(
572         nb_epochs_finished,
573         model,
574         args.result_dir,
575         log_string,
576         args.deterministic_synthesis,
577     )
578
579 for n_epoch in range(nb_epochs_finished, nb_epochs):
580     learning_rate = learning_rate_schedule[n_epoch]
581
582     log_string(f"learning_rate {learning_rate}")
583
584     if args.optim == "sgd":
585         optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
586     elif args.optim == "adam":
587         optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
588     elif args.optim == "adamw":
589         optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
590     else:
591         raise ValueError(f"Unknown optimizer {args.optim}.")
592
593     model.train()
594
595     nb_train_samples, acc_train_loss = 0, 0.0
596
597     for input in task.batches(split="train"):
598         input = input.to(device)
599         output = model(mygpt.BracketedSequence(input)).x
600         loss = F.cross_entropy(output.transpose(1, 2), input)
601         acc_train_loss += loss.item() * input.size(0)
602         nb_train_samples += input.size(0)
603         nb_samples_seen += input.size(0)
604
605         optimizer.zero_grad()
606         loss.backward()
607         optimizer.step()
608
609     with torch.autograd.no_grad():
610         model.eval()
611
612         nb_test_samples, acc_test_loss = 0, 0.0
613
614         for input in task.batches(split="test"):
615             input = input.to(device)
616
617             output = model(mygpt.BracketedSequence(input)).x
618             loss = F.cross_entropy(output.transpose(1, 2), input)
619             acc_test_loss += loss.item() * input.size(0)
620             nb_test_samples += input.size(0)
621
622         train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples))
623         test_perplexity = math.exp(min(100, acc_test_loss / nb_test_samples))
624
625         log_string(
626             f"perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}"
627         )
628
629         task.produce_results(
630             n_epoch, model, args.result_dir, log_string, args.deterministic_synthesis
631         )
632
633     checkpoint = {
634         "nb_epochs_finished": n_epoch + 1,
635         "model_state": model.state_dict(),
636         "rng_state": torch.get_rng_state(),
637     }
638
639     if torch.cuda.is_available():
640         checkpoint["cuda_rng_state"] = torch.cuda.get_rng_state()
641
642     checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name)
643     torch.save(checkpoint, checkpoint_name)
644     log_string(f"saved checkpoint {checkpoint_name}")
645
646 ######################################################################