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

sperate train and val dataset.

This commit is contained in:
Colin 2024-02-26 23:59:00 +08:00
parent d1906629ab
commit 8120be66a6
1 changed files with 5 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
wit/lit_train.py
View File

@ -26,7 +26,6 @@ dataset_name = ["/home/colin/develop/dataset/liwu/MNBVC/wiki"]
dataset_name = ["/home/colin/develop/dataset/liwu/MNBVC/wiki/20230198/58.jsonl.gz"] dataset_name = ["/home/colin/develop/dataset/liwu/MNBVC/wiki/20230198/58.jsonl.gz"]
train_batch_size = 256 train_batch_size = 256
val_batch_size = 16 val_batch_size = 16
limit_val_batches = 128
num_proc = 8 num_proc = 8
max_epochs = 1000 max_epochs = 1000
strategy = "fsdp" strategy = "fsdp"
@ -35,10 +34,10 @@ seed = 42
class SpecialDataset(Dataset): class SpecialDataset(Dataset):
def __init__(self, size=65536): def __init__(self, start, end, size=65536):
self.size = size self.size = size
self.features = [] self.features = []
a = torch.randint(0, 1024, [size]) a = torch.randint(start, end, [size])
self.data = torch.stack([a, a * 2, a * 3, a * 4]).permute(1, 0) self.data = torch.stack([a, a * 2, a * 3, a * 4]).permute(1, 0)
def __len__(self): def __len__(self):
@ -144,8 +143,8 @@ if __name__ == "__main__":
train_dataset = ConcatDataset(train_dataset_list) train_dataset = ConcatDataset(train_dataset_list)
val_dataset = ConcatDataset(val_dataset_list) val_dataset = ConcatDataset(val_dataset_list)
train_dataset = SpecialDataset() train_dataset = SpecialDataset(0, 1000, 65536)
val_dataset = SpecialDataset() val_dataset = SpecialDataset(1000, 1024, 1024)
train_dataloader = DataLoader( train_dataloader = DataLoader(
train_dataset, train_dataset,
@ -166,13 +165,7 @@ if __name__ == "__main__":
torch.set_float32_matmul_precision("medium") torch.set_float32_matmul_precision("medium")
precision = precision precision = precision
lit_trainer = pl.Trainer( lit_trainer = pl.Trainer(accelerator="gpu", precision=precision, strategy=strategy, max_epochs=max_epochs)
accelerator="gpu",
precision=precision,
strategy=strategy,
max_epochs=max_epochs,
limit_val_batches=limit_val_batches,
)
lit_trainer.fit( lit_trainer.fit(
lit_module, lit_module,
train_dataloaders=train_dataloader, train_dataloaders=train_dataloader,