use gloo and update Dataloader

This commit is contained in:
Stardust·减
2023-09-07 20:45:51 +08:00
committed by GitHub
parent 1069066390
commit eb83b8b58b

View File

@@ -49,7 +49,7 @@ global_step = 0
def run(): def run():
dist.init_process_group( dist.init_process_group(
backend="nccl", init_method="env://" backend="gloo", init_method="env://" # Due to some training problem,we proposed to use gloo instead of nccl.
) # Use torchrun instead of mp.spawn ) # Use torchrun instead of mp.spawn
rank = dist.get_rank() rank = dist.get_rank()
n_gpus = dist.get_world_size() n_gpus = dist.get_world_size()
@@ -71,16 +71,19 @@ def run():
num_replicas=n_gpus, num_replicas=n_gpus,
rank=rank, rank=rank,
shuffle=True, shuffle=True,
persistent_workers=True,
prefetch_factor=4
) )
collate_fn = TextAudioSpeakerCollate() collate_fn = TextAudioSpeakerCollate()
train_loader = DataLoader( train_loader = DataLoader(
train_dataset, train_dataset,
num_workers=0, num_workers=16,
shuffle=False, shuffle=False,
pin_memory=True, pin_memory=True,
collate_fn=collate_fn, collate_fn=collate_fn,
batch_sampler=train_sampler, batch_sampler=train_sampler,
) # 0 worker
) # DataLoader config could be adjusted.
if rank == 0: if rank == 0:
eval_dataset = TextAudioSpeakerLoader(hps.data.validation_files, hps.data) eval_dataset = TextAudioSpeakerLoader(hps.data.validation_files, hps.data)
eval_loader = DataLoader( eval_loader = DataLoader(