W&B: fix DDP with wandb disabled (#5163)

utils/loggers/__init__.py

@@ -3,9 +3,11 @@
 Logging utils
 """
 
+import os
 import warnings
 from threading import Thread
 
+import pkg_resources as pkg
 import torch
 from torch.utils.tensorboard import SummaryWriter
 
@@ -15,11 +17,16 @@ from utils.plots import plot_images, plot_results
 from utils.torch_utils import de_parallel
 
 LOGGERS = ('csv', 'tb', 'wandb')  # text-file, TensorBoard, Weights & Biases
+RANK = int(os.getenv('RANK', -1))
 
 try:
     import wandb
 
     assert hasattr(wandb, '__version__')  # verify package import not local dir
+    if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in [0, -1]:
+        wandb_login_success = wandb.login(timeout=30)
+        if not wandb_login_success:
+            wandb = None
 except (ImportError, AssertionError):
     wandb = None
 

utils/loggers/wandb/wandb_utils.py

@@ -20,16 +20,6 @@ from utils.datasets import img2label_paths
 from utils.general import check_dataset, check_file
 
 RANK = int(os.getenv('RANK', -1))
-
-try:
-    import wandb
-
-    assert hasattr(wandb, '__version__')  # verify package import not local dir
-    if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in [0, -1]:
-        wandb.login(timeout=30)
-except (ImportError, AssertionError):
-    wandb = None
-
 WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'