colin
/
Witllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refine train.

This commit is contained in:
Colin 2024-04-19 19:12:38 +08:00
parent 2d415d9e44
commit 24957aa2ae
1 changed files with 9 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
wit/train.py
View File

@ -1,8 +1,3 @@
import argparse
from functools import partial
from itertools import chain
from typing import Dict, Tuple
import pytorch_lightning as pl import pytorch_lightning as pl
import torch import torch
@ -24,23 +19,24 @@ max_epochs = 1000
strategy = "auto" strategy = "auto"
resume_from_ckpt_path = None resume_from_ckpt_path = None
seed = 42 seed = 42
dataloader_works = 2
vocab_size = 256 vocab_size = 256
level_ratio = 6 level_ratio = 5
level = 4 level = 5
dataset_level = 1.5 dataset_level = 1.5
min_subitem = 2 min_subitem = 2
hidden_size = 1024 # 128 1024 2048 32 hidden_size = 128 # 128 1024 2048 32
num_attention_heads = 16 # 8 8 16 num_attention_heads = 16 # 8 8 16
num_hidden_layers = 6 # 6 12 24 3 num_hidden_layers = 6 # 6 12 24 3
mask_level = [0] mask_level = [0, 1]
mask_idx = [-1] mask_idx = [0, -1]
# name = "vocab_ratio_level_data_hidden_head_layer" # name = "vocab_ratio_level_data_hidden_head_layer"
# name = "mask_level_idx" # name = "mask_level_idx"
name = "small" name = "hard"
ver = f"{vocab_size}" + "_" + f"{level_ratio}" + "_" + f"{level}" + "_" + f"{min_subitem}" + "_" + f"{dataset_level}" ver = f"{vocab_size}" + "_" + f"{level_ratio}" + "_" + f"{level}" + "_" + f"{min_subitem}" + "_" + f"{dataset_level}"
ver = ver + "_" + f"{hidden_size}" + "_" + f"{num_attention_heads}" + "_" + f"{num_hidden_layers}" ver = ver + "_" + f"{hidden_size}" + "_" + f"{num_attention_heads}" + "_" + f"{num_hidden_layers}"
@ -63,8 +59,8 @@ if __name__ == "__main__":
raw_dataset = MeaningDataset(start, start + size, size, vocab_size, level_ratio, min_subitem) raw_dataset = MeaningDataset(start, start + size, size, vocab_size, level_ratio, min_subitem)
raw_dataset.set_mask(mask_level, mask_idx) raw_dataset.set_mask(mask_level, mask_idx)
train_dataset, val_dataset = raw_dataset.split(0.9) train_dataset, val_dataset = raw_dataset.split(0.9)
train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size) train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size).dataloader(dataloader_works)
val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size) val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size).dataloader(dataloader_works)
# for i in range(len(train_dataloader)): # for i in range(len(train_dataloader)):
# print(train_dataloader.print_mapping(i)) # print(train_dataloader.print_mapping(i))