github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-12 12:47:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

update markdown docs (english) (#60)

Browse Source
This commit is contained in:
Frank Lee
2021-12-10 14:37:33 +08:00
committed by GitHub
parent da01c234e1
commit 9a0466534c
10 changed files with 341 additions and 374 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
colossalai/utils/data_sampler/data_parallel_sampler.py
View File

@@ -108,7 +108,14 @@ class DataParallelSampler(Sampler):
self.epoch = epoch
def get_dataloader(dataset, shuffle=False, seed=1024, add_sampler=True, **kwargs):
def get_dataloader(dataset,
shuffle=False,
seed=1024,
add_sampler=True,
drop_last=False,
pin_memory=False,
num_workers=0,
**kwargs):
'''Set up a deterministic dataloader (also configure seed workers, samplers and whether shuffle or not)
.. note: when pipeline parallel is enabled, shuffle cannot be True
@@ -141,9 +148,15 @@ def get_dataloader(dataset, shuffle=False, seed=1024, add_sampler=True, **kwargs
return DataLoader(dataset,
worker_init_fn=seed_worker,
shuffle=shuffle,
drop_last=drop_last,
pin_memory=pin_memory,
num_workers=num_workers,
**_kwargs)
else:
return DataLoader(dataset,
sampler=sampler,
worker_init_fn=seed_worker,
drop_last=drop_last,
pin_memory=pin_memory,
num_workers=num_workers,
**_kwargs)