diff --git a/egs/aishell/paraformer/run.sh b/egs/aishell/paraformer/run.sh index 60078a03a..ed6ddfe7c 100755 --- a/egs/aishell/paraformer/run.sh +++ b/egs/aishell/paraformer/run.sh @@ -180,7 +180,6 @@ if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then --input_size $feats_dim \ --ngpu $gpu_num \ --num_worker_count $count \ - --multiprocessing_distributed true \ --dist_init_method $init_method \ --dist_world_size $world_size \ --dist_rank $rank \ diff --git a/funasr/bin/train.py b/funasr/bin/train.py index 474e85753..7a9986cf4 100755 --- a/funasr/bin/train.py +++ b/funasr/bin/train.py @@ -77,6 +77,12 @@ def get_parser(): help="Whether to use the find_unused_parameters in " "torch.nn.parallel.DistributedDataParallel ", ) + parser.add_argument( + "--gpu_id", + type=int, + default=0, + help="local gpu id.", + ) # cudnn related parser.add_argument( @@ -399,6 +405,7 @@ if __name__ == '__main__': torch.backends.cudnn.deterministic = args.cudnn_deterministic # ddp init + os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id) args.distributed = args.dist_world_size > 1 distributed_option = build_distributed(args) diff --git a/funasr/datasets/small_datasets/sequence_iter_factory.py b/funasr/datasets/small_datasets/sequence_iter_factory.py index 707894313..810389e80 100644 --- a/funasr/datasets/small_datasets/sequence_iter_factory.py +++ b/funasr/datasets/small_datasets/sequence_iter_factory.py @@ -62,7 +62,7 @@ class SequenceIterFactory(AbsIterFactory): # sampler dataset_conf = args.dataset_conf batch_sampler = LengthBatchSampler( - batch_bins=dataset_conf["batch_size"], + batch_bins=dataset_conf["batch_size"] * args.ngpu, shape_files=shape_files, sort_in_batch=dataset_conf["sort_in_batch"] if hasattr(dataset_conf, "sort_in_batch") else "descending", sort_batch=dataset_conf["sort_batch"] if hasattr(dataset_conf, "sort_batch") else "ascending",