From 7a6306f1f81e2227f42ef22759289b046384809e Mon Sep 17 00:00:00 2001
From: lizz <innerlee@users.noreply.github.com>
Date: Tue, 11 Jun 2019 15:06:09 +0800
Subject: [PATCH] Set LOCAL_RANK in script (#781)
* Set LOCAL_RANK in script.
Fix multiple checkpoint download bug
Signed-off-by: innerlee <363664470@qq.com>
* beautify
Signed-off-by: innerlee <363664470@qq.com>
---
tools/test.py | 3 +++
tools/train.py | 3 +++
2 files changed, 6 insertions(+)
diff --git a/tools/test.py b/tools/test.py
index ada5e60..af950aa 100644
--- a/tools/test.py
+++ b/tools/test.py
@@ -1,4 +1,5 @@
import argparse
+import os
import os.path as osp
import shutil
import tempfile
@@ -119,6 +120,8 @@ def parse_args():
help='job launcher')
parser.add_argument('--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)
return args
diff --git a/tools/train.py b/tools/train.py
index b8f21d1..d8bb9dc 100644
--- a/tools/train.py
+++ b/tools/train.py
@@ -1,6 +1,7 @@
from __future__ import division
import argparse
+import os
from mmcv import Config
from mmdet import __version__
@@ -35,6 +36,8 @@ def parse_args():
help='job launcher')
parser.add_argument('--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)
return args
--
GitLab