From: François Fleuret Date: Fri, 13 Oct 2023 11:51:34 +0000 (+0200) Subject: Update. X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=commitdiff_plain;ds=inline;h=f44ab6863f93ae348e66ffbf52251d96d3b5453c;hp=cb52b31a10ccf9b8df95114efb6a8039c1e006b6;p=picoclvr.git Update. --- diff --git a/qmlp.y b/qmlp.py similarity index 54% rename from qmlp.y rename to qmlp.py index 7b97edb..e12f0e1 100755 --- a/qmlp.y +++ b/qmlp.py @@ -39,40 +39,24 @@ def dequantize(q, xmin, xmax): ###################################################################### -def create_model(): - hidden_dim = 32 - - model = nn.Sequential( - nn.Linear(2, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, 2), - ) - - return model - - -###################################################################### def generate_sets_and_params( - nb_mlps, + batch_nb_mlps, nb_samples, batch_size, nb_epochs, device=torch.device("cpu"), print_log=False, ): - data_input = torch.zeros(nb_mlps, 2 * nb_samples, 2, device=device) + data_input = torch.zeros(batch_nb_mlps, 2 * nb_samples, 2, device=device) data_targets = torch.zeros( - nb_mlps, 2 * nb_samples, dtype=torch.int64, device=device + batch_nb_mlps, 2 * nb_samples, dtype=torch.int64, device=device ) while (data_targets.float().mean(-1) - 0.5).abs().max() > 0.1: i = (data_targets.float().mean(-1) - 0.5).abs() > 0.1 nb = i.sum() - print(f"{nb=}") nb_rec = 2 support = torch.rand(nb, nb_rec, 2, 3, device=device) * 2 - 1 @@ -108,10 +92,10 @@ def generate_sets_and_params( test_targets = test_targets hidden_dim = 32 - w1 = torch.randn(nb_mlps, hidden_dim, 2, device=device) / math.sqrt(2) - b1 = torch.zeros(nb_mlps, hidden_dim, device=device) - w2 = torch.randn(nb_mlps, 2, hidden_dim, device=device) / math.sqrt(hidden_dim) - b2 = torch.zeros(nb_mlps, 2, device=device) + w1 = torch.randn(batch_nb_mlps, hidden_dim, 2, device=device) / math.sqrt(2) + b1 = torch.zeros(batch_nb_mlps, hidden_dim, device=device) + w2 = torch.randn(batch_nb_mlps, 2, hidden_dim, device=device) / math.sqrt(hidden_dim) + b2 = torch.zeros(batch_nb_mlps, 2, device=device) w1.requires_grad_() b1.requires_grad_() @@ -158,13 +142,13 @@ def generate_sets_and_params( # print(f"{k=} {acc_train_loss=} {train_error=}") q_params = torch.cat( - [quantize(p.view(nb_mlps, -1), -2, 2) for p in [w1, b1, w2, b2]], dim=1 + [quantize(p.view(batch_nb_mlps, -1), -2, 2) for p in [w1, b1, w2, b2]], dim=1 ) q_train_set = torch.cat([q_train_input, train_targets[:, :, None]], -1).reshape( - nb_mlps, -1 + batch_nb_mlps, -1 ) q_test_set = torch.cat([q_test_input, test_targets[:, :, None]], -1).reshape( - nb_mlps, -1 + batch_nb_mlps, -1 ) return q_train_set, q_test_set, q_params @@ -173,51 +157,59 @@ def generate_sets_and_params( ###################################################################### -def evaluate_q_params(q_params, q_set, batch_size=25, device=torch.device("cpu")): +def evaluate_q_params(q_params, q_set, batch_size=25, device=torch.device("cpu"), nb_mlps_per_batch=1024): + + errors = [] nb_mlps = q_params.size(0) - hidden_dim = 32 - w1 = torch.empty(nb_mlps, hidden_dim, 2, device=device) - b1 = torch.empty(nb_mlps, hidden_dim, device=device) - w2 = torch.empty(nb_mlps, 2, hidden_dim, device=device) - b2 = torch.empty(nb_mlps, 2, device=device) - - with torch.no_grad(): - k = 0 - for p in [w1, b1, w2, b2]: - print(f"{p.size()=}") - x = dequantize(q_params[:, k : k + p.numel() // nb_mlps], -2, 2).view( - p.size() - ) - p.copy_(x) - k += p.numel() // nb_mlps - q_set = q_set.view(nb_mlps, -1, 3) - data_input = dequantize(q_set[:, :, :2], -1, 1).to(device) - data_targets = q_set[:, :, 2].to(device) + for n in range(0,nb_mlps,nb_mlps_per_batch): + batch_nb_mlps = min(nb_mlps_per_batch,nb_mlps-n) + batch_q_params = q_params[n:n+batch_nb_mlps] + batch_q_set = q_set[n:n+batch_nb_mlps] + hidden_dim = 32 + w1 = torch.empty(batch_nb_mlps, hidden_dim, 2, device=device) + b1 = torch.empty(batch_nb_mlps, hidden_dim, device=device) + w2 = torch.empty(batch_nb_mlps, 2, hidden_dim, device=device) + b2 = torch.empty(batch_nb_mlps, 2, device=device) - print(f"{data_input.size()=} {data_targets.size()=}") + with torch.no_grad(): + k = 0 + for p in [w1, b1, w2, b2]: + print(f"{p.size()=}") + x = dequantize(batch_q_params[:, k : k + p.numel() // batch_nb_mlps], -2, 2).view( + p.size() + ) + p.copy_(x) + k += p.numel() // batch_nb_mlps - criterion = nn.CrossEntropyLoss() - criterion.to(device) + batch_q_set = batch_q_set.view(batch_nb_mlps, -1, 3) + data_input = dequantize(batch_q_set[:, :, :2], -1, 1).to(device) + data_targets = batch_q_set[:, :, 2].to(device) + + print(f"{data_input.size()=} {data_targets.size()=}") + + criterion = nn.CrossEntropyLoss() + criterion.to(device) + + acc_loss = 0.0 + nb_errors = 0 - acc_loss = 0.0 - nb_errors = 0 + for input, targets in zip( + data_input.split(batch_size, dim=1), data_targets.split(batch_size, dim=1) + ): + h = torch.einsum("mij,mnj->mni", w1, input) + b1[:, None, :] + h = F.relu(h) + output = torch.einsum("mij,mnj->mni", w2, h) + b2[:, None, :] + loss = F.cross_entropy(output.reshape(-1, output.size(-1)), targets.reshape(-1)) + acc_loss += loss.item() * input.size(0) + wta = output.argmax(-1) + nb_errors += (wta != targets).long().sum(-1) - for input, targets in zip( - data_input.split(batch_size, dim=1), data_targets.split(batch_size, dim=1) - ): - h = torch.einsum("mij,mnj->mni", w1, input) + b1[:, None, :] - h = F.relu(h) - output = torch.einsum("mij,mnj->mni", w2, h) + b2[:, None, :] - loss = F.cross_entropy(output.reshape(-1, output.size(-1)), targets.reshape(-1)) - acc_loss += loss.item() * input.size(0) - wta = output.argmax(-1) - nb_errors += (wta != targets).long().sum(-1) + errors.append(nb_errors / data_input.size(1)) + acc_loss = acc_loss / data_input.size(1) - error = nb_errors / data_input.size(1) - acc_loss = acc_loss / data_input.size(1) - return error + return torch.cat(errors) ###################################################################### @@ -229,40 +221,41 @@ def generate_sequence_and_test_set( batch_size, nb_epochs, device, + nb_mlps_per_batch=1024, ): - q_train_set, q_test_set, q_params = generate_sets_and_params( - nb_mlps, - nb_samples, - batch_size, - nb_epochs, - device=device, - ) - input = torch.cat( - [ - q_train_set, - q_train_set.new_full( - ( - q_train_set.size(0), - 1, + inputs, q_test_sets = [],[] + + for n in range(0,nb_mlps,nb_mlps_per_batch): + q_train_set, q_test_set, q_params = generate_sets_and_params( + batch_nb_mlps = min(nb_mlps_per_batch, nb_mlps - n), + nb_samples=nb_samples, + batch_size=batch_size, + nb_epochs=nb_epochs, + device=device, + ) + + inputs.append(torch.cat( + [ + q_train_set, + q_train_set.new_full( + ( + q_train_set.size(0), + 1, + ), + nb_quantization_levels, ), - nb_quantization_levels, - ), - q_params, - ], - dim=-1, - ) + q_params, + ], + dim=-1, + )) - print(f"SANITY #1 {q_train_set.size()=} {q_params.size()=} {input.size()=}") + q_test_sets.append(q_test_set) - ar_mask = ( - (torch.arange(input.size(0), device=input.device) > q_train_set.size(0) + 1) - .long() - .view(1, -1) - .reshape(nb_mlps, -1) - ) + input = torch.cat(inputs) + q_test_set = torch.cat(q_test_sets) - return input, ar_mask, q_test_set + return input, q_test_set ###################################################################### @@ -270,7 +263,7 @@ def generate_sequence_and_test_set( if __name__ == "__main__": import time - nb_mlps, nb_samples = 128, 200 + batch_nb_mlps, nb_samples = 128, 500 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -278,26 +271,22 @@ if __name__ == "__main__": data = [] - for n in range(2): - data.append( - generate_sequence_and_test_set( - nb_mlps=nb_mlps, - nb_samples=nb_samples, - device=device, - batch_size=25, - nb_epochs=250, - ) - ) + input, q_test_set = generate_sequence_and_test_set( + nb_mlps=batch_nb_mlps, + nb_samples=nb_samples, + device=device, + batch_size=25, + nb_epochs=250, + nb_mlps_per_batch=17 + ) end_time = time.perf_counter() - nb = sum([i.size(0) for i, _, _ in data]) - print(f"{nb / (end_time - start_time):.02f} samples per second") - - for input, ar_mask, q_test_set in data: - q_train_set = input[:, : nb_samples * 3] - q_params = input[:, nb_samples * 3 + 1 :] - print(f"SANITY #2 {q_train_set.size()=} {q_params.size()=} {input.size()=}") - error_train = evaluate_q_params(q_params, q_train_set) - print(f"train {error_train*100}%") - error_test = evaluate_q_params(q_params, q_test_set) - print(f"test {error_test*100}%") + print(f"{input.size(0) / (end_time - start_time):.02f} samples per second") + + q_train_set = input[:, : nb_samples * 3] + q_params = input[:, nb_samples * 3 + 1 :] + print(f"SANITY #2 {q_train_set.size()=} {q_params.size()=} {input.size()=}") + error_train = evaluate_q_params(q_params, q_train_set, nb_mlps_per_batch=17) + print(f"train {error_train*100}%") + error_test = evaluate_q_params(q_params, q_test_set, nb_mlps_per_batch=17) + print(f"test {error_test*100}%") diff --git a/tasks.py b/tasks.py index 183c3cf..ea10d7c 100755 --- a/tasks.py +++ b/tasks.py @@ -1550,3 +1550,105 @@ class Grid(Task): ###################################################################### + +import qmlp + + +class QMLP(Task): + + ###################### + + def __init__( + self, + nb_train_samples, + nb_test_samples, + batch_size, + logger=None, + device=torch.device("cpu"), + ): + super().__init__() + + self.device = device + self.batch_size = batch_size + + if logger is not None: + logger( + f"generating {nb_train_samples+nb_test_samples} samples (can take some time)" + ) + + self.train_descr = self.grid_factory.generate_samples( + nb_train_samples, lambda r: tqdm.tqdm(r) + ) + self.test_descr = self.grid_factory.generate_samples( + nb_test_samples, lambda r: tqdm.tqdm(r) + ) + + # Build the tokenizer + tokens = set() + for d in [self.train_descr, self.test_descr]: + for s in d: + for t in s.strip().split(" "): + tokens.add(t) + # make this set a sorted list to get the same tensors given + # the same descr + tokens = list(tokens) + tokens.sort() + tokens = ["#"] + tokens + self.token2id = dict([(t, n) for n, t in enumerate(tokens)]) + self.id2token = dict([(n, t) for n, t in enumerate(tokens)]) + self.t_nul = self.token2id["#"] + self.t_true = self.token2id["true"] + self.t_false = self.token2id["false"] + + # Tokenize the train and test sets + self.train_input = self.str2tensor(self.train_descr) + self.test_input = self.str2tensor(self.test_descr) + + def batches(self, split="train"): + assert split in {"train", "test"} + input = self.train_input if split == "train" else self.test_input + for batch in tqdm.tqdm( + input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}" + ): + yield self.trim(batch) + + def vocabulary_size(self): + return len(self.token2id) + + def produce_results( + self, n_epoch, model, result_dir, logger, deterministic_synthesis + ): + correct = self.test_input[:1000] + result = correct.clone() + ar_mask = torch.logical_or(result == self.t_true, result == self.t_false).long() + result *= 1 - ar_mask # paraaaaanoiaaaaaaa + + logger(f"----------------------------------------------------------") + + for e in self.tensor2str(result[:10]): + logger(f"test_before {e}") + + masked_inplace_autoregression( + model, + self.batch_size, + result, + ar_mask, + deterministic_synthesis, + device=self.device, + ) + + logger(f"----------------------------------------------------------") + + for e in self.tensor2str(result[:10]): + logger(f"test_after {e}") + + logger(f"----------------------------------------------------------") + + nb_total = ar_mask.sum().item() + nb_correct = ((correct == result).long() * ar_mask).sum().item() + + logger(f"test_performance {n_epoch} {nb_total=} {nb_correct=}") + logger(f"main_test_accuracy {n_epoch} {nb_correct / nb_total}") + + +######################################################################