merge:合并

BFDS_train.py

@@ -3,6 +3,19 @@ import logging
 import warnings
 import json
 from datetime import datetime
+import requests
+
+if __name__ == "__main__":
+    try:
+        # 这里尝试连接hugging face连接不上就换国内镜像源
+        response = requests.get("https://huggingface.co", timeout=5)
+        if response.status_code == 200:
+            print("成功连接到 Hugging Face")
+        else:
+            print(f"连接失败，状态码: {response.status_code}")
+    except requests.exceptions.RequestException:
+        os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+        print(f"无法连接到 Hugging Face:换源到{os.environ['HF_ENDPOINT']}")
 
 from utils.logger import setlogger
 from utils.train import train_utils
@@ -28,15 +41,12 @@ class Argument:
         self.model_name = "ResNet_1d"  # 模型名
         self.bottleneck = True  # 是否使用bottleneck层
         self.bottleneck_num = 256  # bottleneck层的输出维数
-        self.pretrained = False  # 是否使用预训练模型
 
         # 训练
         self.batch_size = 64  # 批次大小
         self.cuda_device = "0"  # 训练设备
-        self.last_batch = False  # 是否保留最后的不完整批次
-        self.max_epoch = 10  # 训练最大轮数
+        self.max_epoch = 2  # 训练最大轮数
         self.num_workers = 0  # 训练设备数
-        self.pretrained = False  # 是否加载预训练模型
 
         # 数据记录
         self.checkpoint_dir = "./checkpoint"  # 参数保存路径
@@ -58,7 +68,7 @@ class Argument:
 
         # 基于映射
         self.distance_option = True  # 是否采用基于映射的损失
-        self.distance_loss = "JMMD"  # 损失模型 MK-MMD/JMMD/CORAL
+        self.distance_loss = "MK-MMD"  # 损失模型 MK-MMD/JMMD/CORAL
         self.distance_tradeoff = "Step"  # 损失的trade_off参数 Cons/Step
         self.distance_lambda = 1  # 若调整模式为Cons，指定其具体值
 
@@ -74,31 +84,55 @@ class Argument:
         # 输出可视化
         self.wavelet = "cmor1.5-1.0"  # 小波类型
 
-    def update_param(self, param_name, param_value):
-        if hasattr(self, param_name):
-            setattr(self, param_name, param_value)
-        else:
-            raise AttributeError(f"Parameter '{param_name}' does not exist.")
-
-
-def update_param(args, batch_size, optimizer, learning_rate, scheduler, transfer_method, distance_loss):
-    if transfer_method not in ["基于映射", "基于领域对抗"]:
-        return "错误: 迁移学习方式无效，请选择 '基于映射' 或 '基于领域对抗'。"
-    args.update_param("batch_size", batch_size)
-    args.update_param("opt", optimizer.lower())
-    args.update_param("lr", learning_rate)
-    args.update_param("lr_scheduler", scheduler)
-    if transfer_method == "基于映射":
-        args.update_param("adversarial_option", False)
-        args.update_param("distance_option", True)
-    elif transfer_method == "基于领域对抗":
-        args.update_param("adversarial_option", True)
-        args.update_param("distance_option", False)
-    args.update_param("distance_loss", distance_loss)
-    # 返回所有参数
-    # FIXME __dict__
-    all_params = {attr: getattr(args, attr) for attr in dir(args) if not attr.startswith("__") and not callable(getattr(args, attr))}
-    return json.dumps(all_params, indent=2)
+    def update_params(self, **kwargs):
+        """
+        使用 **kwargs 动态更新 args 的参数。
+        """
+        for param_name, param_value in kwargs.items():
+            if hasattr(self, param_name):
+                setattr(self, param_name, param_value)
+            else:
+                print(f"警告: Parameter '{param_name}' does not exist.")
+
+    def set_recommended_params(self):
+        # 给用户设定的推荐参数
+        recommended_params = {
+            "data_set": "BFDS-Project/Bearing-Fault-Diagnosis-System",
+            "conditions": fetch_all_conditions_from_huggingface("BFDS-Project/Bearing-Fault-Diagnosis-System"),
+            "labels": {"Normal Baseline Data": 0, "Ball": 1, "Inner Race": 2, "Outer Race Centered": 3, "Outer Race Opposite": 4, "Outer Race Orthogonal": 5},
+            "transfer_task": [["CWRU", "CWRU_12k_Drive_End_Bearing_Fault_Data"], ["CWRU", "CWRU_12k_Fan_End_Bearing_Fault_Data"]],
+            "normalize_type": None,
+            "model_name": "CNN",
+            "bottleneck": True,
+            "bottleneck_num": 256,
+            "batch_size": 64,
+            "cuda_device": "0",
+            "max_epoch": 2,
+            "num_workers": 0,
+            "checkpoint_dir": "./checkpoint",
+            "print_step": 50,
+            "opt": "adam",
+            "momentum": 0.9,
+            "weight_decay": 1e-5,
+            "lr": 1e-3,
+            "lr_scheduler": "step",
+            "gamma": 0.1,
+            "steps": [150, 250],
+            "middle_epoch": 0,
+            "distance_option": True,
+            "distance_loss": "JMMD",
+            "distance_tradeoff": "Step",
+            "distance_lambda": 1,
+            "adversarial_option": False,
+            "adversarial_loss": "CDA",
+            "hidden_size": 1024,
+            "grl_option": "Step",
+            "grl_lambda": 1,
+            "adversarial_tradeoff": "Step",
+            "adversarial_lambda": 1,
+            "wavelet": "cmor1.5-1.0",
+        }
+        self.update_params(**recommended_params)
 
 
 if __name__ == "__main__":

BFDS_web.py

@@ -1,11 +1,39 @@
+import os
+import requests
+import zipfile
+
+if __name__ == "__main__":
+    try:
+        # 这里尝试连接hugging face连接不上就换国内镜像源
+        response = requests.get("https://huggingface.co", timeout=5)
+        if response.status_code == 200:
+            print("成功连接到 Hugging Face")
+        else:
+            print(f"连接失败，状态码: {response.status_code}")
+    except requests.exceptions.RequestException:
+        os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+        print(f"无法连接到 Hugging Face:换源到{os.environ['HF_ENDPOINT']}")
+
+
 import gradio as gr
 import matplotlib
 import matplotlib.pyplot as plt
-from BFDS_train import Argument, update_param
-import pandas as pd
+from BFDS_train import Argument
 import torch
 from utils.predict import predict
 
+import logging
+import warnings
+from datetime import datetime
+
+
+from utils.logger import setlogger
+from utils.train import train_utils
+from utils.fetch_conditions import fetch_all_conditions_from_huggingface
+
+dataset_name = "BFDS-Project/Bearing-Fault-Diagnosis-System"
+conditions = fetch_all_conditions_from_huggingface(dataset_name)
+
 # 设置 Matplotlib 的后端为非交互式后端
 matplotlib.use("Agg")
 plt.rcParams.update(
@@ -20,106 +48,359 @@ plt.rcParams.update(
 
 # 初始化 Argument 实例
 args = Argument()
+args.set_recommended_params()
 
 
 # 更新参数的函数
-def transfer_learning(batch_size, optimizer, learning_rate, scheduler, transfer_method, distance_loss):
+def transfer_learning(
+    source_config,
+    source_split,
+    target_path,
+    normalize_type,
+    model_name,
+    bottleneck,
+    bottleneck_num,
+    batch_size,
+    cuda_device,
+    max_epoch,
+    num_workers,
+    opt,
+    momentum,
+    weight_decay,
+    lr,
+    lr_scheduler,
+    gamma,
+    steps_start,
+    steps_end,
+    middle_epoch,
+    distance_option,
+    distance_loss,
+    distance_tradeoff,
+    distance_lambda,
+    adversarial_option,
+    adversarial_loss,
+    hidden_size,
+    grl_option,
+    grl_lambda,
+    adversarial_tradeoff,
+    adversarial_lambda,
+    wavelet,
+):
+    args_params_dict = {
+        "transfer_task": [[source_config, source_split], []],
+        "normalize_type": normalize_type,
+        "model_name": model_name,
+        "bottleneck": bottleneck,
+        "bottleneck_num": bottleneck_num,
+        "batch_size": batch_size,
+        "cuda_device": cuda_device,
+        "max_epoch": max_epoch,
+        "num_workers": num_workers,
+        "opt": opt,
+        "momentum": momentum,
+        "weight_decay": weight_decay,
+        "lr": lr,
+        "lr_scheduler": lr_scheduler,
+        "gamma": gamma,
+        "steps": [steps_start, steps_end],
+        "middle_epoch": middle_epoch,
+        "distance_option": distance_option,
+        "distance_loss": distance_loss,
+        "distance_tradeoff": distance_tradeoff,
+        "distance_lambda": distance_lambda,
+        "adversarial_option": adversarial_option,
+        "adversarial_loss": adversarial_loss,
+        "hidden_size": hidden_size,
+        "grl_option": grl_option,
+        "grl_lambda": grl_lambda,
+        "adversarial_tradeoff": adversarial_tradeoff,
+        "adversarial_lambda": adversarial_lambda,
+        "wavelet": wavelet,
+    }
     # 这里更新参数
-    all_params = update_param(args, batch_size, optimizer, learning_rate, scheduler, transfer_method, distance_loss)
+    if target_path is None:
+        raise ValueError("请上传目标域数据!")
+    args.update_params(**args_params_dict)
     # 这里进行训练
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_device.strip()
+    warnings.filterwarnings("ignore")
+    save_dir = os.path.join(args.checkpoint_dir, args.model_name + "_" + datetime.strftime(datetime.now(), "%m%d-%H%M%S"))
+    setattr(args, "save_dir", save_dir)
+    if not os.path.exists(args.save_dir):
+        os.makedirs(args.save_dir)
+    # 设定日志
+    setlogger(os.path.join(args.save_dir, "train.log"))
+    # 保存超参数
+    for k, v in args.__dict__.items():
+        if k[-3:] != "dir":
+            logging.info(f"{k}: {v}")
+    # 训练
+    trainer = train_utils(args, owned=True, data_path=target_path)
+    trainer.setup()
+    trainer.train()
+    fig = trainer.generate_fig()
 
-    # 这里返回各种结果
-    return all_params
+    # 压缩 save_dir 文件夹
+    zip_filename = f"{trainer.save_dir}.zip"
+    with zipfile.ZipFile(zip_filename, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for root, dirs, files in os.walk(trainer.save_dir):
+            for file in files:
+                zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), os.path.join(trainer.save_dir, "..")))
+
+    return fig, zip_filename
 
 
 # 下面是信号推理的函数
 def signal_inference(model_file, signal_file):
+    result = []
     if model_file is None or signal_file is None:
-        raise ValueError("请上传模型文件和信号数据！")
+        raise ValueError("请上传模型文件和信号数据!")
     model_state_dict = torch.load(model_file)
-    if isinstance(signal_file, list):
-        for signal_file_single in signal_file:
-            signal = pd.read_csv(signal_file_single)
-            # FIXME 最后做成(n,1,128)的形式
-    else:
-        signal = pd.read_csv(signal_file)
-    result = predict(model_state_dict, signal)
+    for signal_file_single in signal_file:
+        result.append(predict(model_state_dict, signal_file_single, args))
     return result
 
 
-# 创建一个绘图函数
-def create_plot():
-    x = [1, 2, 3, 4, 5]
-    y = [1, 4, 9, 16, 25]
-    fig, ax = plt.subplots()
-    ax.plot(x, y, label="y = x^2")
-    ax.set_title("示例折线图")
-    ax.set_xlabel("X 轴")
-    ax.set_ylabel("Y 轴")
-    ax.legend()
-    return fig
+def change_source_split(source_config_radio):
+    source_splits = conditions[source_config_radio]
+    return gr.update(choices=source_splits, value=source_splits[0])
+
+
+def change_bottleneck(bottleneck):
+    return gr.update(visible=bottleneck)
+
+
+def change_opt(opt):
+    if opt == "sgd":
+        return gr.update(visible=True), gr.update(visible=True)
+    elif opt == "adam":
+        return gr.update(visible=False), gr.update(visible=False)
+
+
+def change_lr_scheduler(lr_scheduler):
+    if lr_scheduler == "step":
+        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+    elif lr_scheduler == "exp":
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+    elif lr_scheduler == "stepLR":
+        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+    elif lr_scheduler == "fix":
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+
+def change_steps_start(steps_start, steps_end):
+    if steps_start >= steps_end:
+        steps_start = steps_end - 1
+    return gr.update(value=steps_start, maximum=steps_end - 1)
+
+
+def change_steps_end(steps_start, steps_end):
+    if steps_end <= steps_start:
+        steps_end = steps_start + 1
+    return gr.update(value=steps_end, minimum=steps_start + 1)
 
 
+def change_max_epoch(max_epoch, middle_epoch):
+    if middle_epoch >= max_epoch:
+        middle_epoch = max_epoch - 1
+    return gr.update(value=max_epoch, maximum=max_epoch - 1)
+
+
+def change_middle_epoch(max_epoch, middle_epoch):
+    if middle_epoch >= max_epoch:
+        middle_epoch = max_epoch - 1
+    return gr.update(value=middle_epoch, maximum=max_epoch - 1)
+
+
+def change_distance_option(distance_option, distance_tradeoff):
+    if distance_option:
+        return gr.update(value=False), gr.update(visible=distance_option), gr.update(visible=distance_option), gr.update(visible=(distance_option and distance_tradeoff == "Cons"))
+    else:
+        return gr.update(value=False), gr.update(visible=distance_option), gr.update(visible=distance_option), gr.update(visible=(distance_option and distance_tradeoff == "Cons"))
+
+
+def change_adversarial_option(adversarial_option, adversarial_tradeoff):
+    return (
+        gr.update(value=not adversarial_option),
+        gr.update(visible=adversarial_option),
+        gr.update(visible=adversarial_option),
+        gr.update(visible=adversarial_option),
+        gr.update(visible=adversarial_option),
+        gr.update(visible=adversarial_option),
+        gr.update(visible=(adversarial_option and adversarial_tradeoff == "Cons")),
+    )
+
+
+def change_distance_tradeoff(distance_option, distance_tradeoff):
+    return gr.update(visible=(distance_option and distance_tradeoff == "Cons"))
+
+
+def change_adversarial_tradeoff(adversarial_option, adversarial_tradeoff):
+    return (gr.update(visible=(adversarial_option and adversarial_tradeoff == "Cons")),)
+
+
+with open("docs/BFDS_font.html", "r", encoding="utf-8") as f:
+    BFDS_font_html = f.read()
+
+# gradio BFDS_web.py --demo-name app
 with gr.Blocks(title="BFDS WebUI") as app:
+    gr.HTML(BFDS_font_html)
+    gr.Markdown("""
+    # 轴承故障诊断系统
+    基于深度迁移学习的智能轴承故障诊断系统。支持多种迁移学习算法、信号处理方法和故障诊断模型。
+    """)
     with gr.Tab("模型训练"):
         gr.Markdown("在此模块中，您可以选择不同的迁移学习方法进行模型训练。")
-        with gr.Tab("使用预训练模型"):
-            gr.Markdown("使用预训练模型进行迁移学习，您可以选择以下参数进行配置：")
-            with gr.Row():
-                with gr.Column():
-                    batch_size_slider = gr.Slider(1, 258, label="batch_size", step=1, value=args.batch_size)
-                    optimizer_radio = gr.Radio(
-                        label="选择优化器",
-                        choices=["Adam", "SGD", "RMSprop"],
-                        value=args.opt.capitalize(),
-                    )
-                    learning_rate_slider = gr.Slider(1e-5, 1e-2, label="学习率", step=1e-5, value=args.lr)
-                    scheduler_radio = gr.Radio(
-                        label="学习率调度器",
-                        choices=["step", "exp", "stepLR", "fix"],
-                        value=args.lr_scheduler,
-                    )
-                    transfer_method_radio = gr.Radio(
-                        label="迁移学习方式",
-                        choices=["基于映射", "基于领域对抗"],
-                        value="基于领域对抗" if args.adversarial_option else "基于映射",
-                    )
-                with gr.Column():
-                    distance_loss_radio = gr.Radio(
-                        label="距离损失函数",
-                        choices=["MK-MMD", "JMMD", "CORAL"],
-                        value="MK-MMD",  # 修复默认值为有效选项
-                    )
-            update_button = gr.Button("开始训练")
-            with gr.Row():
-                with gr.Column():
-                    # FIXME 需要弄好看一点
-                    args_all_params = gr.Textbox(label="更新结果", lines=8)
-                with gr.Column():
-                    gr.Plot(create_plot)
-        with gr.Tab("不使用预训练模型"):
-            gr.Markdown("使用从零开始训练的方式，不依赖预训练模型。")
+        with gr.Row():
+            with gr.Column():
+                source_config_radio = gr.Radio(
+                    label="选择源域数据集名称",
+                    choices=list(conditions.keys()),
+                    value=args.transfer_task[0][0],
+                )
+                source_split_radio = gr.Radio(
+                    label="选择源域数据集工况",
+                    choices=conditions[args.transfer_task[0][0]],
+                    value=args.transfer_task[0][1],
+                )
+                target_file = gr.File(label="目标域数据集", file_count="single", file_types=[".csv"])
+                normalize_type_radio = gr.Radio(
+                    label="选择归一化方式",
+                    choices=["mean-std", "min-max", None],
+                    value=args.normalize_type,
+                )
+                model_name_radio = gr.Radio(
+                    label="选择模型名称",
+                    choices=["CNN"],
+                    value=args.model_name,
+                )
+                bottleneck_checkbox = gr.Checkbox(
+                    label="是否使用瓶颈层",
+                    value=args.bottleneck,
+                )
+                bottleneck_num_slider = gr.Slider(1, 1024, label="瓶颈层神经元个数", step=1, value=args.bottleneck_num, visible=args.bottleneck)
+                batch_size_slider = gr.Slider(1, 258, label="batch_size", step=1, value=args.batch_size)
+                cuda_device_radio = gr.Radio(
+                    label="选择GPU设备",
+                    choices=["0"],
+                    value=args.cuda_device,
+                )
+                max_epoch_slider = gr.Slider(args.middle_epoch + 1, 100, label="max_epoch", step=1, value=args.max_epoch)
+                num_workers_slider = gr.Slider(1, 16, label="num_workers", step=1, value=args.num_workers)
+                opt_radio = gr.Radio(
+                    label="选择优化器",
+                    choices=["sgd", "adam"],
+                    value=args.opt,
+                )
+                momentum_slider = gr.Slider(0, 1, label="momentum", step=0.01, value=args.momentum)
+                weight_decay_slider = gr.Slider(1e-5, 1e-1, label="weight_decay", step=1e-5, value=args.weight_decay)
+                lr_slider = gr.Slider(1e-5, 1e-2, label="学习率", step=1e-5, value=args.lr)
+                lr_scheduler_radio = gr.Radio(
+                    label="学习率调度器",
+                    choices=["step", "exp", "stepLR", "fix"],
+                    value=args.lr_scheduler,
+                )
+                gamma_slider = gr.Slider(1e-5, 1e-2, label="gamma", step=1e-5, value=args.gamma, visible=args.lr_scheduler != "fix")
+                steps_start_slider = gr.Slider(1, args.steps[1] - 1, label="steps 第一个值", step=1, value=args.steps[0], visible=(args.lr_scheduler == "step" or args.lr_scheduler == "stepLR"))
+                steps_end_slider = gr.Slider(args.steps[0] + 1, 1000, label="steps 第二个值", step=1, value=args.steps[1], visible=(args.lr_scheduler == "step" or args.lr_scheduler == "stepLR"))
+                middle_epoch_slider = gr.Slider(0, args.max_epoch - 1, label="middle_epoch", step=1, value=args.middle_epoch)
+                wavelet_radio = gr.Radio(
+                    label="选择波形变换",
+                    choices=["cmor1.5-1.0"],
+                    value=args.wavelet,
+                )
+            with gr.Column():
+                #  这两个true和false不能一起出现
+                distance_option_checkbox = gr.Checkbox(
+                    label="是否使用距离损失",
+                    value=args.distance_option,
+                )
+                distance_loss_radio = gr.Radio(label="距离损失函数", choices=["MK-MMD", "JMMD", "CORAL"], value=args.distance_loss, visible=args.distance_option)
+                distance_tradeoff_radio = gr.Radio(label="距离损失权重", choices=["Cons", "Step"], value=args.distance_tradeoff, visible=args.distance_option)
+                distance_lambda_slider = gr.Slider(1, 2, label="距离损失权重", step=1e-5, value=args.distance_lambda, visible=(args.distance_option and args.distance_tradeoff == "Cons"))
+                adversarial_option_checkbox = gr.Checkbox(
+                    label="是否使用对抗损失",
+                    value=args.adversarial_option,
+                )
+                adversarial_loss_radio = gr.Radio(label="对抗损失函数", choices=["DA", "CDA", "CDA+E"], value=args.adversarial_loss, visible=args.adversarial_option)
+                hidden_size_slider = gr.Slider(1, 1024, label="对抗层神经元个数", step=1, value=args.hidden_size, visible=args.adversarial_option)
+                grl_option_radio = gr.Radio(label="是否使用梯度反转层", choices=["Step"], value=args.grl_option, visible=args.adversarial_option)
+                grl_lambda_slider = gr.Slider(1, 2, label="梯度反转层系数", step=1e-5, value=args.grl_lambda, visible=args.adversarial_option)
+                adversarial_tradeoff_radio = gr.Radio(label="对抗损失权重", choices=["Cons", "Step"], value=args.adversarial_tradeoff, visible=args.adversarial_option)
+                adversarial_lambda_slider = gr.Slider(1, 2, label="对抗损失权重", step=1e-5, value=args.adversarial_lambda, visible=(args.adversarial_option and args.adversarial_tradeoff == "Cons"))
+
+        transfer_learning_button = gr.Button("开始训练")
+        with gr.Row():
+            with gr.Column():
+                download_output = gr.File(label="下载训练结果压缩包", interactive=False)
+            with gr.Column():
+                plot_component = gr.Plot(label="训练结果图表")
 
     with gr.Tab("信号推理"):
         model_file = gr.File(label="模型文件", file_count="single", file_types=[".bin", ".pth", ".pt"])
-        with gr.Tab("单次推理"):
-            gr.Markdown("在此模块中，您可以上传信号数据进行推理。")
-            signal_file_single = gr.File(label="上传信号数据", file_count="single", file_types=[".csv"])
-            signal_inference_single_button = gr.Button("开始推理")
-            signal_inference_single_output = gr.Textbox(label="推理结果", lines=8)
-        with gr.Tab("批量推理"):
-            gr.Markdown("在此模块中，您可以上传信号数据进行批量推理。")
-            signal_file_multiple = gr.File(label="上传信号数据", file_count="multiple", file_types=[".csv"])
-            signal_inference_multiple_button = gr.Button("开始批量推理")
-            signal_inference_multiple_output = gr.Textbox(label="批量推理结果", lines=8)
+        gr.Markdown("在此模块中，您可以上传信号数据进行批量推理。")
+        signal_file_multiple = gr.File(label="上传信号数据", file_count="multiple", file_types=[".csv"])
+        signal_inference_button = gr.Button("开始批量推理")
+        signal_inference_output = gr.Textbox(label="批量推理结果", lines=8)
 
     # 下面是所有函数绑定
-    update_button.click(
+    transfer_learning_button.click(
         transfer_learning,
-        inputs=[batch_size_slider, optimizer_radio, learning_rate_slider, scheduler_radio, transfer_method_radio, distance_loss_radio],
-        outputs=args_all_params,
+        inputs=[
+            source_config_radio,
+            source_split_radio,
+            target_file,
+            normalize_type_radio,
+            model_name_radio,
+            bottleneck_checkbox,
+            bottleneck_num_slider,
+            batch_size_slider,
+            cuda_device_radio,
+            max_epoch_slider,
+            num_workers_slider,
+            opt_radio,
+            momentum_slider,
+            weight_decay_slider,
+            lr_slider,
+            lr_scheduler_radio,
+            gamma_slider,
+            steps_start_slider,
+            steps_end_slider,
+            middle_epoch_slider,
+            distance_option_checkbox,
+            distance_loss_radio,
+            distance_tradeoff_radio,
+            distance_lambda_slider,
+            adversarial_option_checkbox,
+            adversarial_loss_radio,
+            hidden_size_slider,
+            grl_option_radio,
+            grl_lambda_slider,
+            adversarial_tradeoff_radio,
+            adversarial_lambda_slider,
+            wavelet_radio,
+        ],
+        outputs=[plot_component, download_output],
+    )
+    source_config_radio.change(change_source_split, inputs=[source_config_radio], outputs=[source_split_radio])
+    opt_radio.change(change_opt, inputs=[opt_radio], outputs=[momentum_slider, weight_decay_slider])
+    bottleneck_checkbox.change(change_bottleneck, inputs=[bottleneck_checkbox], outputs=[bottleneck_num_slider])
+    lr_scheduler_radio.change(change_lr_scheduler, inputs=[lr_scheduler_radio], outputs=[steps_start_slider, steps_end_slider, gamma_slider])
+    steps_start_slider.change(change_steps_start, inputs=[steps_start_slider, steps_end_slider], outputs=[steps_start_slider])
+    steps_end_slider.change(change_steps_end, inputs=[steps_start_slider, steps_end_slider], outputs=[steps_end_slider])
+    max_epoch_slider.change(change_middle_epoch, inputs=[max_epoch_slider, middle_epoch_slider], outputs=[middle_epoch_slider])
+    middle_epoch_slider.change(change_middle_epoch, inputs=[max_epoch_slider, middle_epoch_slider], outputs=[middle_epoch_slider])
+    distance_option_checkbox.change(
+        change_distance_option, inputs=[distance_option_checkbox, distance_tradeoff_radio], outputs=[adversarial_option_checkbox, distance_loss_radio, distance_tradeoff_radio, distance_lambda_slider]
     )
-    signal_inference_single_button.click(signal_inference, inputs=[model_file, signal_file_single], outputs=signal_inference_single_output)
-    signal_inference_multiple_button.click(signal_inference, inputs=[model_file, signal_file_multiple], outputs=signal_inference_multiple_output)
+    adversarial_option_checkbox.change(
+        change_adversarial_option,
+        inputs=[adversarial_option_checkbox, adversarial_tradeoff_radio],
+        outputs=[distance_option_checkbox, adversarial_loss_radio, hidden_size_slider, grl_option_radio, grl_lambda_slider, adversarial_tradeoff_radio, adversarial_lambda_slider],
+    )
+    distance_tradeoff_radio.change(change_distance_tradeoff, inputs=[distance_option_checkbox, distance_tradeoff_radio], outputs=[distance_lambda_slider])
+    adversarial_tradeoff_radio.change(change_adversarial_tradeoff, inputs=[adversarial_option_checkbox, adversarial_tradeoff_radio], outputs=[adversarial_lambda_slider])
+    signal_inference_button.click(signal_inference, inputs=[model_file, signal_file_multiple], outputs=signal_inference_output)
+
 app.queue()
 app.launch()

dataset/dataset.py

@@ -1,4 +1,5 @@
 import pandas as pd
+import numpy as np
 from datasets import load_dataset
 import torch
 from torch.utils.data import Dataset, DataLoader, random_split
@@ -6,11 +7,27 @@ from typing import Optional, Literal
 
 
 def get_dataset(data_set, subset, split):
-    # TODO 换源
     ds = load_dataset(data_set, subset)
     return ds[split].to_pandas()
 
 
+def get_owned_dataset(data_path):
+    # 提供更多读取方式，和预测一起整理一下
+    df = pd.read_csv(data_path).dropna()
+    data = df.values
+    if data.size % 224 != 0:
+        raise ValueError(f"数据大小 {data.size} 不能被 224 整除，无法重塑为 (-1, 224)")
+    # 重塑数据为 (-1, 224)
+    reshaped_data = data.reshape(-1, 224)
+    # 创建一个全是 0 的列
+    zero_column = np.zeros((reshaped_data.shape[0], 1))
+    # 将 reshaped_data 和 zero_column 拼接成新的数组
+    new_data = np.hstack((reshaped_data, zero_column))
+    # 将新的数组转换为 DataFrame
+    owned_df = pd.DataFrame(new_data, columns=[f"col_{i}" for i in range(225)])
+    return owned_df
+
+
 class SignalDataset(Dataset):
     def __init__(self, data_frame: pd.DataFrame, normalize_type: Optional[Literal["mean-std", "min-max"]] = None):
         if normalize_type == "mean-std":
@@ -36,34 +53,38 @@ class SignalDatasetCreator:
         self.source = transfer_task[0]
         self.target = transfer_task[1]
 
-    def data_split(self, batch_size, num_workers, device, transfer_learning=True):
-        if transfer_learning:
-            # get source train and val
-            data_frame = get_dataset(self.data_set, self.source[0], self.source[1])
-            data_set = SignalDataset(data_frame)
-            lengths = [round(0.8 * len(data_set)), len(data_set) - round(0.8 * len(data_set))]
-            train_data, eval_data = random_split(data_set, lengths)
-            source_train = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=(device == "cuda"),drop_last=True)
-            source_val = DataLoader(dataset=eval_data, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=(device == "cuda"),drop_last=True)
-            # get target train and val
-            data_frame=get_dataset(self.data_set, self.target[0], self.target[1])
-            data_set = SignalDataset(data_frame)
-            lengths = [round(0.8 * len(data_set)), len(data_set) - round(0.8 * len(data_set))]
-            train_data, eval_data = random_split(data_set, lengths)
-            target_train = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=(device == "cuda"),drop_last=True)
-            target_val = DataLoader(dataset=eval_data, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=(device == "cuda"),drop_last=True)
-            return source_train, source_val, target_train, target_val
-        else:
-            # get source train and val
-            data_frame = get_dataset(self.data_set, self.source[0], self.source[1])
-            data_set = SignalDataset(data_frame)
-            lengths = [round(0.8 * len(data_set)), len(data_set) - round(0.8 * len(data_set))]
-            train_data, eval_data = random_split(data_set, lengths)
-            source_train = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=(device == "cuda"))
-            source_val = DataLoader(dataset=eval_data, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=(device == "cuda"))
-            # get target val
-            get_dataset(self.data_set, self.target[0], self.target[1])
-            data_frame = get_dataset(self.data_set, self.target[0, 0], self.target[0, 1])
-            data_set = SignalDataset(data_frame)
-            target_val = DataLoader(dataset=data_set, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=(device == "cuda"))
-            return source_train, source_val, target_val
+    def data_split(self, batch_size, num_workers, device):
+        # 这里源域和目标域都是我们提供用来验证迁移学习的正确性
+        # get source train and val
+        data_frame_source = get_dataset(self.data_set, self.source[0], self.source[1])
+        data_set_source = SignalDataset(data_frame_source)
+        lengths_source = [round(0.8 * len(data_set_source)), len(data_set_source) - round(0.8 * len(data_set_source))]
+        train_data_source, eval_data_source = random_split(data_set_source, lengths_source)
+        source_train = DataLoader(dataset=train_data_source, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=(device == "cuda"), drop_last=True)
+        source_val = DataLoader(dataset=eval_data_source, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=(device == "cuda"), drop_last=True)
+        # get target train and val
+        data_frame_target = get_dataset(self.data_set, self.target[0], self.target[1])
+        data_set_target = SignalDataset(data_frame_target)
+        lengths_target = [round(0.8 * len(data_set_target)), len(data_set_target) - round(0.8 * len(data_set_target))]
+        train_data_target, eval_data_target = random_split(data_set_target, lengths_target)
+        target_train = DataLoader(dataset=train_data_target, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=(device == "cuda"), drop_last=True)
+        target_val = DataLoader(dataset=eval_data_target, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=(device == "cuda"), drop_last=True)
+        return source_train, source_val, target_train, target_val
+
+    def owned_data_split(self, data_path, batch_size, num_workers, device):
+        # 这里目标域是用户自己提供的数据集
+        # get source train and val
+        data_frame_source = get_dataset(self.data_set, self.source[0], self.source[1])
+        data_set_source = SignalDataset(data_frame_source)
+        lengths_source = [round(0.8 * len(data_set_source)), len(data_set_source) - round(0.8 * len(data_set_source))]
+        train_data_source, eval_data_source = random_split(data_set_source, lengths_source)
+        source_train = DataLoader(dataset=train_data_source, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=(device == "cuda"), drop_last=True)
+        source_val = DataLoader(dataset=eval_data_source, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=(device == "cuda"), drop_last=True)
+        # get target train and val
+        data_frame_target = get_owned_dataset(data_path)
+        data_set_target = SignalDataset(data_frame_target)
+        lengths_target = [round(0.8 * len(data_set_target)), len(data_set_target) - round(0.8 * len(data_set_target))]
+        train_data_target, eval_data_target = random_split(data_set_target, lengths_target)
+        target_train = DataLoader(dataset=train_data_target, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=(device == "cuda"), drop_last=True)
+        target_val = DataLoader(dataset=eval_data_target, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=(device == "cuda"), drop_last=True)
+        return source_train, source_val, target_train, target_val

docs/BFDS_font.html

@@ -0,0 +1,31 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Bearing Fault Diagnosis ASCII Art</title>
+    <style>
+        #taag_output_text {
+            font-family: "Courier New", ui-monospace, monospace;
+            font-size: 10pt;
+            white-space: pre;
+            overflow-wrap: break-word;
+            margin-top: 15px;
+            margin-bottom: 15px;
+            float: left;
+        }
+    </style>
+</head>
+<body>
+    <pre id="taag_output_text">
+ ________  ________ ________  ________                 ________  ________  ________        ___  _______   ________ _________   
+|\   __  \|\  _____\\   ___ \|\   ____\               |\   __  \|\   __  \|\   __  \      |\  \|\  ___ \ |\   ____\\___   ___\ 
+\ \  \|\ /\ \  \__/\ \  \_|\ \ \  \___|_  ____________\ \  \|\  \ \  \|\  \ \  \|\  \     \ \  \ \   __/|\ \  \___\|___ \  \_| 
+ \ \   __  \ \   __\\ \  \ \\ \ \_____  \|\____________\ \   ____\ \   _  _\ \  \\\  \  __ \ \  \ \  \_|/_\ \  \       \ \  \  
+  \ \  \|\  \ \  \_| \ \  \_\\ \|____|\  \|____________|\ \  \___|\ \  \\  \\ \  \\\  \|\  \\_\  \ \  \_|\ \ \  \____   \ \  \ 
+   \ \_______\ \__\   \ \_______\____\_\  \              \ \__\    \ \__\\ _\\ \_______\ \________\ \_______\ \_______\  \ \__\
+    \|_______|\|__|    \|_______|\_________\              \|__|     \|__|\|__|\|_______|\|________|\|_______|\|_______|   \|__|
+                                 \|_________|                                                                                   
+    </pre>
+</body>
+</html>

utils/fetch_conditions.py

@@ -1,9 +1,23 @@
+import os
+import requests
+
+if __name__ == "__main__":
+    try:
+        # 这里尝试连接hugging face连接不上就换国内镜像源
+        response = requests.get("https://huggingface.co", timeout=5)
+        if response.status_code == 200:
+            print("成功连接到 Hugging Face")
+        else:
+            print(f"连接失败，状态码: {response.status_code}")
+    except requests.exceptions.RequestException:
+        os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+        print(f"无法连接到 Hugging Face:换源到{os.environ['HF_ENDPOINT']}")
+
 from datasets import get_dataset_config_names, get_dataset_split_names
 import json
 
 
 def fetch_all_conditions_from_huggingface(dataset_name):
-    # TODO 换源
     """所有数据集的subset和split
     具体见网页https://huggingface.co/datasets/BFDS-Project/Bearing-Fault-Diagnosis-System
     Args:
@@ -26,4 +40,6 @@ if __name__ == "__main__":
     dataset_name = "BFDS-Project/Bearing-Fault-Diagnosis-System"
     conditions = fetch_all_conditions_from_huggingface(dataset_name)
     print("huggingface上的数据集配置和分割信息:")
-    print(json.dumps(conditions, indent=2))
+    # print(json.dumps(conditions, indent=2))
+    # 返回conditions的key用数组存储
+    print(conditions[0][0])

utils/logger.py

@@ -8,6 +8,9 @@ def setlogger(path):
         path(_str_): log文件保存路径
     """
     logger = logging.getLogger()
+    if logger.hasHandlers():  # 检查是否已经有处理器
+        logger.handlers.clear()  # 清除已有的处理器，避免重复输出
+
     logger.setLevel(logging.INFO)
     logFormatter = logging.Formatter("%(asctime)s %(message)s", "%m-%d %H:%M:%S")  # 格式为  月-日 时：分：秒
 

utils/predict.py

@@ -1,28 +1,45 @@
-from models.CNN import cnn_features
 import torch
 import torch.nn as nn
+import pandas as pd
+import librosa
+from pathlib import Path
+import models
 
-from main import Argument
 
-# FIXME 这里的 Argument 类看是直接导入还是后期需要main传进来
-args = Argument()
+def audio_to_signal(audio_file, sr=None):
+    signal, _ = librosa.load(audio_file, sr=sr)
+    return signal
 
 
-def predict(model_state_dict, signal):
-    model = cnn_features()
+def csv_to_signal(signal_file):
+    signal = pd.read_csv(signal_file).to_numpy().flatten()
+    return signal
+
+
+# 修改backbone
+def predict(model_state_dict, signal_file, args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = getattr(models, args.model_name)().to(device)
     bottleneck_layer = nn.Sequential(
         nn.Linear(model.output_num(), args.bottleneck_num),
         nn.ReLU(inplace=True),
         nn.Dropout(),
-    )
-    classifier_layer = nn.Linear(args.bottleneck_num, 10)
-    model_all = nn.Sequential(model, bottleneck_layer, classifier_layer)
+    ).to(device)
+    classifier_layer = nn.Linear(args.bottleneck_num, len(args.labels)).to(device)
+    model_all = nn.Sequential(model, bottleneck_layer, classifier_layer).to(device)
     model_all.load_state_dict(model_state_dict)
-
-    # 设置为评估模式
+    # 模型预测
     model_all.eval()
-    # 进行预测
     with torch.no_grad():
-        # FIXME signal的处理还没有写
-        output = model_all(torch.randn(10, 1, 1024))
-    return torch.argmax(output, dim=1)
+        # 根据文件后缀选择处理方式
+        file_extension = Path(signal_file).suffix
+        if file_extension == ".csv":
+            signal = csv_to_signal(signal_file).reshape(-1, 1, 224)
+        elif file_extension in [".wav", ".mp3"]:
+            signal = audio_to_signal(signal_file).reshape(-1, 1, 224)
+        else:
+            raise ValueError(f"Unsupported file type: {file_extension}")
+        signal = torch.tensor(signal, dtype=torch.float32).to(device)
+        output = model_all(signal)
+        predictions = output.mean(dim=0)
+    return predictions

utils/train.py

@@ -14,12 +14,14 @@ import matplotlib.pyplot as plt
 import models
 from models.AdversarialNet import AdversarialNet, calc_coeff, grl_hook, Entropy
 from dataset.dataset import SignalDatasetCreator
-from .loss import DAN, JAN, CORAL
+from utils.loss import DAN, JAN, CORAL
 
 
 class train_utils:
-    def __init__(self, args):
+    def __init__(self, args, owned=False, data_path=None):
         self.args = args
+        self.owned = owned
+        self.data_path = data_path
 
     def setup(self):
         args = self.args
@@ -38,17 +40,18 @@ class train_utils:
             logging.info(f"using {self.device_count} cpu")
 
         # 加载数据集
-        signal_dataset_creator = SignalDatasetCreator(args.data_set, args.labels, args.transfer_task)
-        self.datasets = {}
-        self.datasets["source_train"], self.datasets["source_val"], self.datasets["target_train"], self.datasets["target_val"] = signal_dataset_creator.data_split(
-            args.batch_size, args.num_workers, self.device, transfer_learning=True
-        )
-        self.dataloaders = {
-            "source_train": self.datasets["source_train"],
-            "source_val": self.datasets["source_val"],
-            "target_train": self.datasets["target_train"],
-            "target_val": self.datasets["target_val"],
-        }
+        if self.owned:
+            signal_dataset_creator = SignalDatasetCreator(args.data_set, args.labels, args.transfer_task)
+            self.dataloaders = {}
+            self.dataloaders["source_train"], self.dataloaders["source_val"], self.dataloaders["target_train"], self.dataloaders["target_val"] = signal_dataset_creator.owned_data_split(
+                self.data_path, args.batch_size, args.num_workers, self.device
+            )
+        else:
+            signal_dataset_creator = SignalDatasetCreator(args.data_set, args.labels, args.transfer_task)
+            self.dataloaders = {}
+            self.dataloaders["source_train"], self.dataloaders["source_val"], self.dataloaders["target_train"], self.dataloaders["target_val"] = signal_dataset_creator.data_split(
+                args.batch_size, args.num_workers, self.device
+            )
         # 定义模型
         self.model = getattr(models, args.model_name)()
         if args.bottleneck:
@@ -91,9 +94,7 @@ class train_utils:
                     )
             else:
                 if args.bottleneck_num:
-                    self.AdversarialNet = AdversarialNet(
-                        in_feature=args.bottleneck_num, hidden_size=args.hidden_size, max_iter=self.max_iter, grl_option=args.grl_option, grl_lambda=args.grl_lambda
-                    )
+                    self.AdversarialNet = AdversarialNet(in_feature=args.bottleneck_num, hidden_size=args.hidden_size, max_iter=self.max_iter, grl_option=args.grl_option, grl_lambda=args.grl_lambda)
                 else:
                     self.AdversarialNet = AdversarialNet(
                         in_feature=self.model.output_num(), hidden_size=args.hidden_size, max_iter=self.max_iter, grl_option=args.grl_option, grl_lambda=args.grl_lambda
@@ -336,9 +337,7 @@ class train_utils:
                                     domain_label_source = torch.zeros(labels.size(0)).float()
                                     domain_label_target = torch.ones(inputs.size(0) - labels.size(0)).float()
                                     adversarial_label = torch.cat((domain_label_source, domain_label_target), dim=0).to(self.device)
-                                    weight = torch.cat(
-                                        (entropy_source / torch.sum(entropy_source).detach().item(), entropy_target / torch.sum(entropy_target).detach().item()), dim=0
-                                    )
+                                    weight = torch.cat((entropy_source / torch.sum(entropy_source).detach().item(), entropy_target / torch.sum(entropy_target).detach().item()), dim=0)
 
                                     # 展开权重，对损失重新加权
                                     adversarial_loss = torch.sum(weight.view(-1, 1) * self.adversarial_loss(adversarial_out.squeeze(), adversarial_label))
@@ -451,3 +450,27 @@ class train_utils:
 
         plt.tight_layout()
         plt.show()
+
+    def generate_fig(self):
+        args = self.args
+
+        fig, axs = plt.subplots(1, 2, figsize=(14, 6))
+
+        axs[0].set_title("Accuracy")
+        axs[0].set_xlabel("epoches")
+        axs[0].set_ylabel("accuracy")
+        axs[0].plot(range(args.max_epoch), self.acc["source_train"], label="source_train")
+        axs[0].plot(range(args.max_epoch), self.acc["source_val"], label="source_val")
+        axs[0].plot(range(args.max_epoch), self.acc["target_val"], label="target_val")
+        axs[0].legend()
+
+        axs[1].set_title(f"Loss Function: {args.distance_loss}")
+        axs[1].set_xlabel("epoches")
+        axs[1].set_ylabel("loss")
+        axs[1].plot(range(args.max_epoch), self.loss["source_train"], label="source_train")
+        axs[1].plot(range(args.max_epoch), self.loss["source_val"], label="source_val")
+        axs[1].plot(range(args.max_epoch), self.loss["target_val"], label="target_val")
+        axs[1].legend()
+
+        plt.tight_layout()
+        return fig