Feat/Converter-Wandb

swanlab/cli/main.py

@@ -147,15 +147,11 @@ def login(api_key: str, relogin: bool, **kwargs):
 
 # ---------------------------------- 转换命令，用于转换其他实验跟踪工具 ----------------------------------
 @cli.command()
-@click.argument(
-    "convert_dir",
-    type=str,
-)
 @click.option(
     "--type",
     "-t",
     default="tensorboard",
-    type=click.Choice(["tensorboard"]),
+    type=click.Choice(["tensorboard", "wandb"]),
     help="The type of the experiment tracking tool you want to convert to.",
 )
 @click.option(
@@ -184,19 +180,68 @@ def login(api_key: str, relogin: bool, **kwargs):
     type=str,
     help="The directory where the swanlab log files are stored.",
 )
-def convert(convert_dir: str, type: str, project: str, cloud: bool, workspace: str, logdir: str, **kwargs):
+@click.option(
+    "--tb_logdir",
+    type=str,
+    help="The directory where the tensorboard log files are stored.",
+)
+@click.option(
+    "--wb-project",
+    type=str,
+    help="The project name of the wandb runs.",
+)
+@click.option(
+    "--wb-entity",
+    type=str,
+    help="The entity name of the wandb runs.",
+)
+@click.option(
+    "--wb-runid",
+    type=str,
+    help="The run_id of the wandb run.",
+)
+def convert(
+    type: str,
+    project: str,
+    cloud: bool,
+    workspace: str,
+    logdir: str,
+    tb_logdir: str,
+    wb_project: str,
+    wb_entity: str,
+    wb_runid: str,
+    **kwargs,
+):
     """Convert the log files of other experiment tracking tools to SwanLab."""
     if type == "tensorboard":
         from swanlab.converter import TFBConverter
 
         tfb_converter = TFBConverter(
-            convert_dir=convert_dir,
+            convert_dir=tb_logdir,
             project=project,
             workspace=workspace,
             cloud=cloud,
             logdir=logdir,
         )
         tfb_converter.run()
+
+    elif type == "wandb":
+        from swanlab.converter import WandbConverter
+
+        print(wb_project, wb_entity, wb_runid)
+
+        wb_converter = WandbConverter(
+            project=project,
+            workspace=workspace,
+            cloud=cloud,
+            logdir=logdir,
+        )
+        wb_converter.run(
+            wb_project=wb_project,
+            wb_entity=wb_entity,
+            wb_run_id=wb_runid,
+        )
+
     else:
         raise ValueError("The type of the experiment tracking tool you want to convert to is not supported.")
 

swanlab/converter/__init__.py

@@ -1 +1,2 @@
 from .tfb import TFBConverter
+from .wb import WandbConverter

swanlab/converter/tfb/tfb_converter.py

@@ -1,7 +1,3 @@
-"""
-ISSUE: https://github.com/SwanHubX/SwanLab/issues/437
-"""
-
 import os
 import swanlab
 from datetime import datetime
@@ -15,7 +11,6 @@ def __init__(
         convert_dir: str,
         project: str = None,
         workspace: str = None,
-        config: dict = None,
         cloud: bool = True,
         logdir: str = None,
         **kwargs,
@@ -24,7 +19,6 @@ def __init__(
         self.project = project
         self.workspace = workspace
         self.cloud = cloud
-        self.config = config
         self.logdir = logdir
 
     def run(self, depth=3):
@@ -61,9 +55,6 @@ def run(self, depth=3):
                         logdir=self.logdir,
                     )
 
-                    if self.config:
-                        run.config.update(self.config)
-
                     """
                     根据tag提取数据, 格式为{tag: [(step, value, wall_time), ...]}, example:
                     data_by_tags = {
@@ -100,7 +91,7 @@ def run(self, depth=3):
 
                     # 计算完整的运行时间
                     runtime = max(times) - min(times)
-                    swanlab.config.update({"RunTime": runtime})
+                    swanlab.config.update({"RunTime(s)": runtime})
 
                     # 结束当前实验
                     run.finish()

swanlab/converter/wb/__init__.py

@@ -0,0 +1 @@
+from .wb_converter import WandbConverter

swanlab/converter/wb/wb_converter.py

@@ -0,0 +1,97 @@
+"""
+------example.py------
+from swanlab.converter import WandbConverter
+
+wb_converter = WandbConverter()
+wb_converter.run(wb_project="WANDB_PROJECT_NAME", wb_entity="WANDB_USERNAME")
+"""
+
+import swanlab
+from swanlab.log import swanlog as swl
+
+
+class WandbConverter:
+    def __init__(
+        self,
+        project: str = None,
+        workspace: str = None,
+        cloud: bool = True,
+        logdir: str = None,
+        **kwargs,
+    ):
+        self.project = project
+        self.workspace = workspace
+        self.cloud = cloud
+        self.logdir = logdir
+
+    def parse_wandb_logs(self, wb_project: str, wb_entity: str, wb_run_id: str = None):
+        try:
+            import wandb
+        except ImportError as e:
+            raise TypeError(
+                "Wandb Converter requires wandb when process tfevents file. Install with 'pip install wandb'."
+            )
+
+        client = wandb.Api()
+
+        if wb_run_id is None:
+            # process all runs
+            runs = client.runs(wb_entity + "/" + wb_project)
+        else:
+            # get the run by run_id
+            run = client.run(f"{wb_entity}/{wb_project}/{wb_run_id}")
+            runs = (run,)
+
+        for iter, wb_run in enumerate(runs):
+            swl.info(f"Conversion progress: {iter+1}/{len(runs)}")
+
+            if swanlab.get_run() is None:
+                swanlab_run = swanlab.init(
+                    project=wb_project if self.project is None else self.project,
+                    workspace=self.workspace,
+                    experiment_name=wb_run.name,
+                    description=wb_run.notes,
+                    cloud=self.cloud,
+                    logdir=self.logdir,
+                )
+            else:
+                swanlab_run = swanlab.get_run()
+
+            wb_config = {
+                "wandb_run_id": wb_run.id,
+                "wandb_run_name": wb_run.name,
+                "Created Time": wb_run.created_at,
+                "wandb_user": wb_run.user,
+                "wandb_tags": wb_run.tags,
+                "wandb_url": wb_run.url,
+                "wandb_metadata": wb_run.metadata,
+            }
+
+            swanlab_run.config.update(wb_config)
+            swanlab_run.config.update(wb_run.config)
+
+            keys = [key for key in wb_run.history(stream="default").keys() if not key.startswith("_")]
+
+            # 记录标量指标
+            for record in wb_run.scan_history():
+                step = record.get("_step")
+                for key in keys:
+                    value = record.get(key)
+                    # 如果value是None或者是dict类型，则跳过
+                    if value is None or isinstance(value, dict) or not isinstance(value, (float, int)):
+                        # 如果是多媒体数据，如图像，value是这个格式
+                        # image {'format': 'png', 'path': 'media/images/image_13_3bbb7517118b6af0307c.png', 'sha256': '3bbb7517118b6af0307cbe1b26f6d94b68797874112de716df4b5b50e01ddc24', 'size': 30168, 'height': 100, 'width': 100, '_type': 'image-file'}
+                        continue
+                    swanlab_run.log({key: value}, step=step)
+
+            # 结束此轮实验
+            swanlab_run.finish()
+
+    def run(self, wb_project: str, wb_entity: str, wb_run_id: str = None):
+        swl.info("Start converting Wandb Runs to SwanLab...")
+        self.parse_wandb_logs(
+            wb_project=wb_project,
+            wb_entity=wb_entity,
+            wb_run_id=wb_run_id,
+        )
+        swl.info("Finished converting Wandb Runs to SwanLab.")