X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=main.py;h=6c1def76800aba41b7462d0ec977461562207d67;hb=0dbca4cef7405fb92689e5d2542f1d4761d658a3;hp=83c305bd351d5f4c2ec4659679d2a54a1001ba61;hpb=3b4983952f4b5089528783347a4d0d76c7fb4fc8;p=mygpt.git diff --git a/main.py b/main.py index 83c305b..6c1def7 100755 --- a/main.py +++ b/main.py @@ -69,12 +69,24 @@ parser.add_argument('--dropout', parser.add_argument('--synthesis_sampling', action='store_true', default = True) +parser.add_argument('--no_checkpoint', + action='store_true', default = False) + parser.add_argument('--checkpoint_name', type = str, default = 'checkpoint.pth') +############################## +# picoclvr options + parser.add_argument('--picoclvr_many_colors', action='store_true', default = False) +parser.add_argument('--picoclvr_height', + type = int, default = 12) + +parser.add_argument('--picoclvr_width', + type = int, default = 16) + ###################################################################### args = parser.parse_args() @@ -118,37 +130,43 @@ import picoclvr class TaskPicoCLVR(Task): def __init__(self, batch_size, - height = 6, width = 8, many_colors = False, + height, width, many_colors = False, device = torch.device('cpu')): + def generate_descr(nb): + descr = picoclvr.generate( + nb, + height = self.height, width = self.width, + many_colors = many_colors + ) + + descr = [ s.strip().split(' ') for s in descr ] + l = max([ len(s) for s in descr ]) + descr = [ s + [ '' ] * (l - len(s)) for s in descr ] + + return descr + + self.height = height + self.width = width self.batch_size = batch_size self.device = device nb = args.data_size if args.data_size > 0 else 250000 - descr = picoclvr.generate( - nb, - height = height, width = width, - many_colors = many_colors - ) - - # self.test_descr = descr[:nb // 5] - # self.train_descr = descr[nb // 5:] - - descr = [ s.strip().split(' ') for s in descr ] - l = max([ len(s) for s in descr ]) - descr = [ s + [ '' ] * (l - len(s)) for s in descr ] + self.train_descr = generate_descr((nb * 4) // 5) + self.test_descr = generate_descr((nb * 1) // 5) + # Build the tokenizer tokens = set() - for s in descr: - for t in s: tokens.add(t) + for d in [ self.train_descr, self.test_descr ]: + for s in d: + for t in s: tokens.add(t) self.token2id = dict([ (t, n) for n, t in enumerate(tokens) ]) self.id2token = dict([ (n, t) for n, t in enumerate(tokens) ]) - t = [ [ self.token2id[u] for u in s ] for s in descr ] - data_input = torch.tensor(t, device = self.device) - - self.test_input = data_input[:nb // 5] - self.train_input = data_input[nb // 5:] + t = [ [ self.token2id[u] for u in s ] for s in self.train_descr ] + self.train_input = torch.tensor(t, device = self.device) + t = [ [ self.token2id[u] for u in s ] for s in self.test_descr ] + self.test_input = torch.tensor(t, device = self.device) def batches(self, split = 'train'): assert split in { 'train', 'test' } @@ -180,7 +198,9 @@ class TaskPicoCLVR(Task): return ' '.join(t_primer + t_generated) - def produce_results(self, n_epoch, model, nb_tokens = 50): + def produce_results(self, n_epoch, model, nb_tokens = None): + if nb_tokens is None: + nb_tokens = self.height * self.width + 3 descr = [ ] nb_per_primer = 8 @@ -194,14 +214,23 @@ class TaskPicoCLVR(Task): for k in range(nb_per_primer): descr.append(self.generate(primer, model, nb_tokens)) - img = [ picoclvr.descr2img(d) for d in descr ] + img = [ picoclvr.descr2img(d, height = self.height, width = self.width) for d in descr ] img = torch.cat(img, 0) - file_name = f'result_picoclvr_{n_epoch:04d}.png' - torchvision.utils.save_image(img / 255., - file_name, nrow = nb_per_primer, pad_value = 0.8) - log_string(f'wrote {file_name}') + image_name = f'result_picoclvr_{n_epoch:04d}.png' + torchvision.utils.save_image( + img / 255., + image_name, nrow = nb_per_primer, pad_value = 0.8 + ) + log_string(f'wrote {image_name}') + + nb_missing = sum( [ + x[2] for x in picoclvr.nb_missing_properties( + descr, + height = self.height, width = self.width + ) + ] ) - log_string(f'nb_misssing {picoclvr.nb_missing_properties(descr)}') + log_string(f'nb_missing {nb_missing / len(descr):.02f}') ###################################################################### @@ -364,7 +393,11 @@ if args.data == 'wiki103': elif args.data == 'mnist': task = TaskMNIST(batch_size = args.batch_size, device = device) elif args.data == 'picoclvr': - task = TaskPicoCLVR(batch_size = args.batch_size, many_colors = args.picoclvr_many_colors, device = device) + task = TaskPicoCLVR(batch_size = args.batch_size, + height = args.picoclvr_height, + width = args.picoclvr_width, + many_colors = args.picoclvr_many_colors, + device = device) else: raise ValueError(f'Unknown dataset {args.data}.') @@ -400,19 +433,23 @@ else: nb_epochs_finished = 0 -try: - checkpoint = torch.load(args.checkpoint_name, map_location = device) - nb_epochs_finished = checkpoint['nb_epochs_finished'] - model.load_state_dict(checkpoint['model_state']) - optimizer.load_state_dict(checkpoint['optimizer_state']) - print(f'Checkpoint loaded with {nb_epochs_finished} epochs finished.') +if args.no_checkpoint: + log_string(f'Not trying to load checkpoint.') -except FileNotFoundError: - print('Starting from scratch.') - -except: - print('Error when loading the checkpoint.') - exit(1) +else: + try: + checkpoint = torch.load(args.checkpoint_name, map_location = device) + nb_epochs_finished = checkpoint['nb_epochs_finished'] + model.load_state_dict(checkpoint['model_state']) + optimizer.load_state_dict(checkpoint['optimizer_state']) + log_string(f'Checkpoint loaded with {nb_epochs_finished} epochs finished.') + + except FileNotFoundError: + log_string('Starting from scratch.') + + except: + log_string('Error when loading the checkpoint.') + exit(1) ######################################################################