mirrors
/
mmocr
mirror of https://github.com/open-mmlab/mmocr.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[Enhancement] Accepts local-rank in train.py and test.py (#1806) 

* [Enhancement] Accepts local-rank

* add

* update
pull/1811/head
Tong Gao 2023-03-27 10:34:54 +08:00 committed by GitHub
parent f47cff5199
commit 73df26d749
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tools/test.py
View File

@ -45,7 +45,10 @@ def parse_args():
help='Job launcher')
parser.add_argument(
'--tta', action='store_true', help='Test time augmentation')
parser.add_argument('--local_rank', type=int, default=0)
# When using PyTorch version >= 2.0.0, the `torch.distributed.launch`
# will pass the `--local-rank` parameter to `tools/test.py` instead
# of `--local_rank`.
parser.add_argument('--local_rank', '--local-rank', type=int, default=0)
args = parser.parse_args()
if 'LOCAL_RANK' not in os.environ:
os.environ['LOCAL_RANK'] = str(args.local_rank)

5
tools/train.py
View File

@ -41,7 +41,10 @@ def parse_args():
choices=['none', 'pytorch', 'slurm', 'mpi'],
default='none',
help='Job launcher')
parser.add_argument('--local_rank', type=int, default=0)
# When using PyTorch version >= 2.0.0, the `torch.distributed.launch`
# will pass the `--local-rank` parameter to `tools/train.py` instead
# of `--local_rank`.
parser.add_argument('--local_rank', '--local-rank', type=int, default=0)
args = parser.parse_args()
if 'LOCAL_RANK' not in os.environ:
os.environ['LOCAL_RANK'] = str(args.local_rank)