36 lines
909 B
Python
36 lines
909 B
Python
|
import os
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
# 路径配置
|
|||
|
|
DATA_ROOT = "./data" # 数据根目录，包含000,001等子文件夹
|
|||
|
|
OUTPUT_DIR = "./output"
|
|||
|
|
MODEL_SAVE_DIR = os.path.join(OUTPUT_DIR, "models")
|
|||
|
|
|
|||
|
|
# 确保目录存在
|
|||
|
|
os.makedirs(MODEL_SAVE_DIR, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 数据配置
|
|||
|
|
IMG_SIZE = 224 # 调整图像大小
|
|||
|
|
BATCH_SIZE = 32
|
|||
|
|
NUM_WORKERS = 4
|
|||
|
|
TRAIN_RATIO = 0.8
|
|||
|
|
VAL_RATIO = 0.2
|
|||
|
|
|
|||
|
|
# 样本数量控制
|
|||
|
|
MAX_SAMPLES_PER_CLASS = 1000 # 每个类别最多读取的样本数
|
|||
|
|
NORMAL_CLASS = "000" # 正常类别的文件夹名
|
|||
|
|
ABNORMAL_CLASSES = ["001", "010", "011", "100", "101", "110", "111"] # 异常类别的文件夹名
|
|||
|
|
|
|||
|
|
# 模型配置
|
|||
|
|
HIDDEN_DIM = 768
|
|||
|
|
NUM_HEADS = 12
|
|||
|
|
NUM_LAYERS = 6
|
|||
|
|
DROPOUT = 0.1
|
|||
|
|
NUM_CLASSES = 2 # 二分类：正常 vs 异常
|
|||
|
|
|
|||
|
|
# 训练配置
|
|||
|
|
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|||
|
|
LEARNING_RATE = 1e-4
|
|||
|
|
WEIGHT_DECAY = 1e-5
|
|||
|
|
NUM_EPOCHS = 50
|
|||
|
|
EARLY_STOPPING_PATIENCE = 10
|