diff --git a/wit/lit_module.py b/wit/lit_module.py index ced90cb..af2fc4d 100644 --- a/wit/lit_module.py +++ b/wit/lit_module.py @@ -93,5 +93,6 @@ class LitModule(pl.LightningModule): stopping_threshold=1, ) lr_monitor = pl.callbacks.LearningRateMonitor(logging_interval="step") - return [checkpoint_callback, lr_monitor] + return [lr_monitor] + # return [checkpoint_callback, lr_monitor] # return [checkpoint_callback, early_stop_callback] diff --git a/wit/meaning_dataset.py b/wit/meaning_dataset.py index 40d8b87..9d52122 100644 --- a/wit/meaning_dataset.py +++ b/wit/meaning_dataset.py @@ -125,9 +125,12 @@ class MeaningDataset(Dataset): self.length.append(len(sq)) unique, counts = np.unique(self.length, return_counts=True) - print("MeaningDataset size: " + str(len(self.length))) - print("MeaningDataset max sequence length: " + str(max(unique))) - print("MeaningDataset most popular sequence length: " + str(unique[np.argmax(counts)])) + print("----------------------------------------------------------------") + print("MeaningDataset start:" + str(start) + " end:" + str(end) + " space:" + str(end - start)) + print("MeaningDataset size:" + str(len(self.length))) + print("MeaningDataset max sequence length:" + str(max(unique))) + print("MeaningDataset most popular sequence length:" + str(unique[np.argmax(counts)])) + print("----------------------------------------------------------------") def __len__(self): return len(self.data) @@ -197,6 +200,8 @@ class BatchGroupMeaningDataloader(Dataset): np.random.shuffle(index_shuffle) index = index[index_shuffle] self.indexBatch = index + print("Dataloader batch size:" + str(batch_size) + " count:" + str(len(index))) + print("Dataloader total:" + str(len(length)) + " drop:" + str(len(length) - len(index) * batch_size)) def __len__(self): return len(self.indexBatch) diff --git a/wit/train.py b/wit/train.py index 79d958b..955c13d 100644 --- a/wit/train.py +++ b/wit/train.py @@ -17,34 +17,26 @@ pretrain_model_name = None # "qwen/Qwen-1_8B-Chat" learning_rate = 0.0001 use_tril_attention_mask = None precision = "32-true" # "precision:bf16-mixed,16-mixed,32-true" -train_batch_size = 32 -val_batch_size = 4 +train_batch_size = 4 +val_batch_size = 1 num_proc = 8 max_epochs = 1000 strategy = "auto" resume_from_ckpt_path = None seed = 42 -vocab_size = 2048 +vocab_size = 1024 level_ratio = 4 level = 4 +dataset_level = 1 hidden_size = 256 # 128 1024 2048 32 num_attention_heads = 8 # 8 8 16 -num_hidden_layers = 1 # 6 12 24 3 +num_hidden_layers = 2 # 6 12 24 3 -name = "vocab_level_hidden_head_layer" -version = ( - str(vocab_size) - + "_" - + str(level_ratio) - + "_" - + str(hidden_size) - + "_" - + str(num_attention_heads) - + "_" - + str(num_hidden_layers) -) +name = "vocab_ratio_level_data_hidden_head_layer" +ver = f"{vocab_size}" + "_" + f"{level_ratio}" + "_" + f"{level}" + "_" + f"{dataset_level}" +ver = ver + "_" + f"{hidden_size}" + "_" + f"{num_attention_heads}" + "_" + f"{num_hidden_layers}" if __name__ == "__main__": torch.manual_seed(seed) @@ -60,9 +52,9 @@ if __name__ == "__main__": start = vocab_size * (level_ratio**level) end = start * level_ratio - size = vocab_size * (level_ratio ** (level / 2)) + size = int(vocab_size * (level_ratio**dataset_level)) raw_dataset = MeaningDataset(start, end, size, vocab_size, level_ratio) - train_dataset, val_dataset = raw_dataset.Split(0.95) + train_dataset, val_dataset = raw_dataset.Split(0.9) train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size) val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size) # it = iter(train_dataloader) @@ -75,7 +67,7 @@ if __name__ == "__main__": accelerator="cuda", devices=[0, 1], precision=precision, - logger=TBLogger("./log/", name=name, version=version, default_hp_metric=False), + logger=TBLogger("./log/", name=name, version=ver, default_hp_metric=False), strategy=strategy, max_epochs=max_epochs, )