Fix missing work_dir

open-mmlab · Oct 22, 2024 · 1ff0afc · 1ff0afc
1 parent c9b5996
commit 1ff0afc
Showing 1 changed file with 14 additions and 6 deletions.
diff --git a/mmengine/visualization/vis_backend.py b/mmengine/visualization/vis_backend.py
@@ -824,12 +824,20 @@ def close(self) -> None:
             return
 
         file_paths = dict()
-        for filename in scandir(self.cfg.work_dir, self._artifact_suffix,
-                                True):
-            file_path = osp.join(self.cfg.work_dir, filename)
-            relative_path = os.path.relpath(file_path, self.cfg.work_dir)
-            dir_path = os.path.dirname(relative_path)
-            file_paths[file_path] = dir_path
+        if (hasattr(self, 'cfg')
+                and osp.isdir(getattr(self.cfg, 'work_dir', ''))):
+            for filename in scandir(
+                    self.cfg.work_dir,
+                    self._artifact_suffix,
+                    recursive=True,
+            ):
+                file_path = str(osp.join(self.cfg.work_dir, filename))
+                relative_path = os.path.relpath(file_path, self.cfg.work_dir)
+                dir_path = os.path.dirname(relative_path)
+                file_paths[file_path] = dir_path
+        else:
+            warnings.warn('self.cfg.work_dir is not set, thus some '
+                          'artifacts will not be logged')
 
         for file_path, dir_path in file_paths.items():
             self._mlflow.log_artifact(file_path, dir_path)