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

Add multithread dataloder support.

This commit is contained in:
Colin 2024-04-17 23:59:09 +08:00
parent 17de117bda
commit a524d01ac3
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
wit/meaning_dataset.py
View File

@ -447,6 +447,14 @@ class BatchGroupMeaningDataloader(Dataset):
s += "--------------------------------------------------------\n" s += "--------------------------------------------------------\n"
return s return s
def detection_collate(batch):
return batch[0]
def dataloader(self, num_workers=1):
return DataLoader(
self, batch_size=1, num_workers=num_workers, collate_fn=BatchGroupMeaningDataloader.detection_collate
)
if __name__ == "__main__": if __name__ == "__main__":