X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=sidebyside;f=main.py;h=aa1b51799a159b2f176154d63cce5893eabd38ad;hb=62533ba50393866c15b322074cad836684dd69e7;hp=ac1e2e3f103abdf6fcc0fb5d736fe1dd3025e80a;hpb=dfeb9072208095669528fc5ae2dedf78f089d9ad;p=mygpt.git diff --git a/main.py b/main.py index ac1e2e3..aa1b517 100755 --- a/main.py +++ b/main.py @@ -18,15 +18,11 @@ import mygpt device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') ###################################################################### - parser = argparse.ArgumentParser(description = 'My own GPT.') parser.add_argument('--log_filename', type = str, default = 'train.log') -parser.add_argument('--download', - action='store_true', default = False) - parser.add_argument('--seed', type = int, default = 0) @@ -46,7 +42,10 @@ parser.add_argument('--optim', type = str, default = 'adam') parser.add_argument('--learning_rate', - type = float, default = 1e-4) + type = float, default = 1e-3) + +parser.add_argument('--learning_rate_end', + type = float, default = 1e-6) parser.add_argument('--dim_model', type = int, default = 512) @@ -66,8 +65,8 @@ parser.add_argument('--nb_blocks', parser.add_argument('--dropout', type = float, default = 0.1) -parser.add_argument('--synthesis_sampling', - action='store_true', default = True) +parser.add_argument('--deterministic_synthesis', + action='store_true', default = False) parser.add_argument('--no_checkpoint', action='store_true', default = False) @@ -114,8 +113,8 @@ for n in vars(args): ###################################################################### def autoregression( - model, - nb_samples, nb_tokens_to_generate, starting_input = None, + model, batch_size, + nb_samples, nb_tokens_to_generate, primer = None, device = torch.device('cpu') ): results = torch.zeros( @@ -123,21 +122,21 @@ def autoregression( dtype = torch.int64, device = device ) - if starting_input is None: + if primer is None: first = 0 else: - first = starting_input.size(1) - results = torch.cat((starting_input, results), 1) + first = primer.size(1) + results = torch.cat((primer, results), 1) - for input in results.split(args.batch_size): - for s in tqdm.tqdm(range(first, input.size(1)), desc = 'synth'): + for input in results.split(batch_size): + for s in range(first, input.size(1)): output = model(input) logits = output[:, s] - if args.synthesis_sampling: + if args.deterministic_synthesis: + t_next = logits.argmax(1) + else: dist = torch.distributions.categorical.Categorical(logits = logits) t_next = dist.sample() - else: - t_next = logits.argmax(1) input[:, s] = t_next return results @@ -151,7 +150,7 @@ class Task: def vocabulary_size(self): pass - def produce_results(self, n_epoch, model, nb_tokens = 50): + def produce_results(self, n_epoch, model): pass ###################################################################### @@ -160,109 +159,141 @@ import picoclvr class TaskPicoCLVR(Task): + # Make a tensor from a list of strings + def tensorize(self, descr): + token_descr = [ s.strip().split(' ') for s in descr ] + l = max([ len(s) for s in token_descr ]) + #token_descr = [ [ '' ] * (l - len(s)) + s for s in token_descr ] + token_descr = [ s + [ '' ] * (l - len(s)) for s in token_descr ] + id_descr = [ [ self.token2id[u] for u in s ] for s in token_descr ] + return torch.tensor(id_descr, device = self.device) + + def trim(self, x, token = ''): + n = self.token2id[token] + i = (1 - (F.pad(x, (1, 1), value = n) == n).min(0).values.long()).cumsum(0) + a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min() + return x[:, a:b] + def __init__(self, batch_size, height, width, nb_colors = 5, device = torch.device('cpu')): def generate_descr(nb): - descr = picoclvr.generate( + return picoclvr.generate( nb, height = self.height, width = self.width, nb_colors = nb_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 + log_string(f'generating {nb} samples (can take some time)') self.train_descr = generate_descr((nb * 4) // 5) self.test_descr = generate_descr((nb * 1) // 5) # Build the tokenizer - tokens = set() + tokens = { '' } for d in [ self.train_descr, self.test_descr ]: for s in d: - for t in s: tokens.add(t) + for t in s.strip().split(' '): 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 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) + # Tokenize the train and test sets + self.train_input = self.tensorize(self.train_descr) + self.test_input = self.tensorize(self.test_descr) def batches(self, split = 'train'): assert split in { 'train', 'test' } - if split == 'train': - for batch in tqdm.tqdm(self.train_input.split(self.batch_size), desc = f'epoch-{split}'): - yield batch - else: - for batch in tqdm.tqdm(self.test_input.split(self.batch_size), desc = f'epoch-{split}'): - yield batch + input = self.train_input if split == 'train' else self.test_input + for batch in tqdm.tqdm(input.split(self.batch_size), desc = f'epoch-{split}'): + yield self.trim(batch) def vocabulary_size(self): return len(self.token2id) - def generate(self, primer, model, nb_tokens): - t_primer = primer.strip().split(' ') - t_generated = [ ] + def test_model(self, n_epoch, model, primers_descr, nb_per_primer=1, generate_images=False): + nb_tokens_to_generate = self.height * self.width + 3 + result_descr = [ ] - for j in range(nb_tokens): - t = [ [ self.token2id[u] for u in t_primer + t_generated ] ] - input = torch.tensor(t, device = self.device) - input = F.pad(input, (0, 1)) # Add the next token, the one to predict - output = model(input) - logits = output[0, -1] - if args.synthesis_sampling: - dist = torch.distributions.categorical.Categorical(logits = logits) - t_next = dist.sample() - else: - t_next = logits.argmax() - t_generated.append(self.id2token[t_next.item()]) - - return ' '.join(t_primer + t_generated) - - 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 - - for primer in [ - 'red above green green top blue right of red ', - 'there is red there is yellow there is blue ', - 'red below yellow yellow below green green below blue red right yellow left green right blue left ', - 'green bottom yellow bottom green left of blue yellow right of blue blue top ', - ]: - - for k in range(nb_per_primer): - descr.append(self.generate(primer, model, nb_tokens)) - - img = [ picoclvr.descr2img(d, height = self.height, width = self.width) for d in descr ] - img = torch.cat(img, 0) - 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 + for primer_descr in primers_descr: + + results = autoregression( + model, + self.batch_size, + nb_samples = nb_per_primer, + nb_tokens_to_generate = nb_tokens_to_generate, + primer = self.tensorize([ primer_descr ]).expand(nb_per_primer, -1), + device = self.device + ) + + l = [ ' '.join([ self.id2token[t.item()] for t in r ]) for r in results ] + result_descr += l + + np = picoclvr.nb_properties( + result_descr, + height = self.height, width = self.width ) - 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 + nb_requested_properties, _, nb_missing_properties = zip(*np) + + log_string(f'nb_requested_properties {sum(nb_requested_properties) / len(result_descr):.02f} nb_missing_properties {sum(nb_missing_properties) / len(result_descr):.02f}') + + np=torch.tensor(np) + count=torch.empty(np[:,0].max()+1,np[:,2].max()+1,dtype=torch.int64) + for i in range(count.size(0)): + for j in range(count.size(1)): + count[i,j]=((np[:,0]==i).long()*(np[:,2]==j).long()).sum() + + if generate_images: + img = [ + picoclvr.descr2img(d, height = self.height, width = self.width) + for d in result_descr + ] + + img = torch.cat(img, 0) + 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}') + + return count + + def produce_results(self, n_epoch, model): + primers_descr = [ + 'red above green green top blue right of red ', + 'there is red there is yellow there is blue ', + 'red below yellow yellow below green green below blue red right yellow left green right blue left ', + 'green bottom yellow bottom green left of blue yellow right of blue blue top ', + ] + + self.test_model( + n_epoch, model, + primers_descr, + nb_per_primer=8, generate_images=True + ) - log_string(f'nb_missing {nb_missing / len(descr):.02f}') + # FAR TOO SLOW!!! + + # test_primers_descr=[ s.split('')[0] for s in self.test_descr ] + + # count=self.test_model( + # n_epoch, model, + # test_primers_descr, + # nb_per_primer=1, generate_images=False + # ) + + # with open(f'perf_{n_epoch:04d}.txt', 'w') as f: + # for i in range(count.size(0)): + # for j in range(count.size(1)): + # f.write(f'{count[i,j]}') + # f.write(" " if j', '' ], + specials = [ '', '' ], min_freq = self.min_freq ) self.vocab.set_default_index(self.vocab[ '' ]) + # makes a tensor from a list of list of tokens def tensorize(self, s): a = max(len(x) for x in s) - return torch.tensor([ self.vocab(x + [ '' ] * (a - len(x))) for x in s ]) + return torch.tensor([ self.vocab(x + [ '' ] * (a - len(x))) for x in s ]) def yield_batches(self, ds): s = [ ] @@ -325,7 +357,8 @@ class TaskWiki103(Task): def vocabulary_size(self): return len(self.vocab) - def produce_results(self, n_epoch, model, nb_tokens = 50): + def produce_results(self, n_epoch, model): + nb_tokens = 50 file_name = f'result_wiki103_{n_epoch:04d}.txt' with open(file_name, 'w') as outfile: @@ -349,13 +382,13 @@ class TaskWiki103(Task): input = F.pad(input, (0, 1)) # Add the next token, the one to predict output = model(input) logits = output[0, -1] - if args.synthesis_sampling: + if args.deterministic_synthesis: + t_next = logits.argmax() + else: dist = torch.distributions.categorical.Categorical(logits = logits) t_next = dist.sample() - else: - t_next = logits.argmax() t_generated.append(self.vocab.lookup_token(t_next)) - if t_generated[-1] == '': break + if t_generated[-1] == '': break s = ' '.join(t_generated) @@ -386,8 +419,9 @@ class TaskMNIST(Task): def vocabulary_size(self): return 256 - def produce_results(self, n_epoch, model, nb_samples = 64): - results = autoregression(model, nb_samples, 28 * 28, device = self.device) + def produce_results(self, n_epoch, model): + nb_samples = 64 + results = autoregression(model, self.batch_size, nb_samples, 28 * 28, device = self.device) image_name = f'result_mnist_{n_epoch:04d}.png' torchvision.utils.save_image(1 - results.reshape(-1, 1, 28, 28) / 255., image_name, nrow = 16, pad_value = 0.8) @@ -432,17 +466,6 @@ log_string(f'nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)') ###################################################################### -if args.optim == 'sgd': - optimizer = torch.optim.SGD(model.parameters(), lr = args.learning_rate) -elif args.optim == 'adam': - optimizer = torch.optim.Adam(model.parameters(), lr = args.learning_rate) -elif args.optim == 'adamw': - optimizer = torch.optim.AdamW(model.parameters(), lr = args.learning_rate) -else: - raise ValueError(f'Unknown optimizer {args.optim}.') - -###################################################################### - nb_epochs_finished = 0 if args.no_checkpoint: @@ -450,10 +473,12 @@ if args.no_checkpoint: else: try: - checkpoint = torch.load(args.checkpoint_name, map_location = device) + checkpoint = torch.load(args.checkpoint_name) nb_epochs_finished = checkpoint['nb_epochs_finished'] model.load_state_dict(checkpoint['model_state']) - optimizer.load_state_dict(checkpoint['optimizer_state']) + torch.set_rng_state(checkpoint['rng_state']) + if torch.cuda.is_available(): + torch.cuda.set_rng_state(checkpoint['cuda_rng_state']) log_string(f'checkpoint loaded with {nb_epochs_finished} epochs finished.') except FileNotFoundError: @@ -471,11 +496,27 @@ token_count = 0 for input in task.batches(split = 'train'): token_count += F.one_hot(input, num_classes = task.vocabulary_size()).sum((0, 1)) token_probas = token_count / token_count.sum() -h = -torch.xlogy(token_probas, token_probas).sum() -train_set_perplexity = math.exp(h) -log_string(f'train set perplexity {train_set_perplexity}') +entropy = -torch.xlogy(token_probas, token_probas).sum() +train_set_perplexity = math.exp(entropy) -for k in range(nb_epochs_finished, nb_epochs): +for n_epoch in range(nb_epochs_finished, nb_epochs): + + if args.learning_rate_end < 0: + lr = args.learning_rate + else: + u = n_epoch / (nb_epochs - 1) + lr = math.exp((1 - u) * math.log(args.learning_rate) + + u * math.log(args.learning_rate_end)) + log_string(f'learning_rate {lr}') + + if args.optim == 'sgd': + optimizer = torch.optim.SGD(model.parameters(), lr = lr) + elif args.optim == 'adam': + optimizer = torch.optim.Adam(model.parameters(), lr = lr) + elif args.optim == 'adamw': + optimizer = torch.optim.AdamW(model.parameters(), lr = lr) + else: + raise ValueError(f'Unknown optimizer {args.optim}.') model.train() @@ -508,16 +549,19 @@ for k in range(nb_epochs_finished, nb_epochs): train_perplexity = math.exp(min(100, acc_train_loss/nb_train_samples)) test_perplexity = math.exp(min(100, acc_test_loss/nb_test_samples)) - log_string(f'perplexity {k} train {train_perplexity} test {test_perplexity}') + log_string(f'perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}') - task.produce_results(k, model) + task.produce_results(n_epoch, model) checkpoint = { - 'nb_epochs_finished': k + 1, + 'nb_epochs_finished': n_epoch + 1, 'model_state': model.state_dict(), - 'optimizer_state': optimizer.state_dict() + 'rng_state': torch.get_rng_state(), } + if torch.cuda.is_available(): + checkpoint['cuda_rng_state'] = torch.cuda.get_rng_state() + torch.save(checkpoint, args.checkpoint_name) ######################################################################