170 lines
7.8 KiB
Python
170 lines
7.8 KiB
Python
# -*- coding:utf-8 -*-
|
||
import configparser
|
||
from daemon import Daemon
|
||
import sys
|
||
from torch.cuda import is_available
|
||
from loguru import logger
|
||
from datetime import datetime
|
||
from utils import create_all_dataloader, train_and_test_model
|
||
import method
|
||
import gc
|
||
import traceback
|
||
import pandas
|
||
import os
|
||
|
||
|
||
class Config(configparser.ConfigParser):
|
||
def __init__(self, defaults=None):
|
||
configparser.ConfigParser.__init__(self, defaults=defaults)
|
||
|
||
def optionxform(self, optionstr):
|
||
return optionstr
|
||
|
||
def as_dict(self):
|
||
d = dict(self._sections)
|
||
for k in d:
|
||
d[k] = dict(d[k])
|
||
return d
|
||
|
||
|
||
class Main(Daemon):
|
||
def __init__(self, pidfile):
|
||
super(Main, self).__init__(pidfile=pidfile)
|
||
current = datetime.now()
|
||
self.start_time = current.strftime("%Y-%m-%d_%H-%M-%S")
|
||
|
||
if len(sys.argv) == 1 or sys.argv[1] == "start":
|
||
self.run()
|
||
elif sys.argv[1] == "stop":
|
||
self.stop()
|
||
elif sys.argv[1] == "daemon":
|
||
self.start()
|
||
else:
|
||
print("Input format error. Please input: python3 main.py start|stop|daemon")
|
||
sys.exit(0)
|
||
|
||
def run(self):
|
||
# 读取配置文件参数
|
||
cf = Config()
|
||
cf.read("./config.ini", encoding='utf8')
|
||
cf_dict = cf.as_dict()
|
||
# 读取数据集名称
|
||
dataset_names = cf.get("Dataset", "name")
|
||
datasets = dataset_names.split(",")
|
||
datasets = [name.strip() for name in datasets]
|
||
# 读取模型名称
|
||
model_names = cf.get("Method", "name")
|
||
models = model_names.split(",")
|
||
models = [name.strip() for name in models]
|
||
# 读取预处理方法
|
||
preprocess_name = cf.get("Preprocess", "name")
|
||
# 读取评估方法
|
||
evaluation_names = cf.get("Evaluation", "name")
|
||
evaluations = evaluation_names.split(",")
|
||
evaluations = [name.strip() for name in evaluations]
|
||
# 读取模型参数文件路径
|
||
model_path = cf_dict["ModelPath"]
|
||
# 读取训练参数
|
||
train = cf.getboolean("BaseParameters", "train")
|
||
epochs = cf.getint("BaseParameters", "epochs")
|
||
batch_size = cf.getint("BaseParameters", "batch_size")
|
||
learning_rate = cf.getfloat("BaseParameters", "learning_rate")
|
||
device = cf.get("BaseParameters", "device")
|
||
if device == "auto":
|
||
device = 'cuda:0' if is_available() else 'cpu'
|
||
# 读取自定义参数
|
||
customs = cf_dict["CustomParameters"]
|
||
|
||
# 建立本次实验记录的路径
|
||
os.makedirs(f"./records", exist_ok=True)
|
||
os.makedirs(f"./records/{self.start_time}", exist_ok=True)
|
||
os.makedirs(f"./records/{self.start_time}/detection_result", exist_ok=True)
|
||
os.makedirs(f"./records/{self.start_time}/model", exist_ok=True)
|
||
|
||
# 初始化日志
|
||
logger.add(f"./records/{self.start_time}/log",
|
||
level='DEBUG',
|
||
format='{time:YYYY-MM-DD HH:mm:ss} - {level} - {file} - {line} - {message}',
|
||
rotation="100 MB")
|
||
|
||
# 核心程序
|
||
self.core(models, datasets, preprocess_name, evaluations, model_path, train, epochs, batch_size, learning_rate,
|
||
device, customs)
|
||
logger.info(f"实验结束，关闭进程")
|
||
|
||
def core(self, models: [str], datasets: [str], preprocess_name: str, evaluations: [str], model_path: {}, train: bool,
|
||
epochs: int, batch_size: int, learning_rate: float, device: str, customs: {}):
|
||
"""
|
||
初始化数据集与模型，并开始训练与测试
|
||
:param models: 训练的模型名称，可包含多个
|
||
:param datasets: 使用的数据集名称，可包含多个
|
||
:param preprocess_name: 预处理方法名称
|
||
:param evaluations: 评估方法名称，可包含多个
|
||
:param model_path: 需要加载模型参数的路径，可包含多个
|
||
:param train: 是否训练，如果为False，则仅测试模型
|
||
:param epochs: 总训练轮数
|
||
:param batch_size: batch的尺寸
|
||
:param learning_rate: 学习率
|
||
:param device: 设备
|
||
:param customs: 自定义参数
|
||
"""
|
||
logger.info(f"加载数据集")
|
||
try:
|
||
# 初始化所有数据集
|
||
all_dataloader = create_all_dataloader(datasets=datasets, input_size=int(customs["input_size"]),
|
||
output_size=int(customs["output_size"]), step=int(customs["step"]),
|
||
batch_size=batch_size, time_index=customs["time_index"] == "true",
|
||
del_column_name=customs["del_column_name"] == "true",
|
||
preprocess_name=preprocess_name)
|
||
except RuntimeError:
|
||
logger.error(traceback.format_exc())
|
||
return
|
||
|
||
# 开始训练与测试
|
||
|
||
for model_name in models:
|
||
try:
|
||
logger.info(f"------------华丽丽的分界线：{model_name} 实验开始------------")
|
||
for i in range(len(all_dataloader)):
|
||
dataloader = all_dataloader[i]
|
||
all_score = {}
|
||
for sub_dataloader in dataloader:
|
||
dataset_name = sub_dataloader["dataset_name"]
|
||
normal_dataloader = sub_dataloader["normal"]
|
||
attack_dataloader = sub_dataloader["attack"]
|
||
logger.info(f"初始化模型 {model_name}")
|
||
model = eval(f"method.{model_name}.Model")(customs=customs, dataloader=normal_dataloader)
|
||
model = model.to(device)
|
||
logger.info(f"模型初始化完成")
|
||
pth_name = model_path[f"{model_name}_{dataset_name}"] if f"{model_name}_{dataset_name}" \
|
||
in model_path else None
|
||
best_score, best_detection = train_and_test_model(start_time=self.start_time, epochs=epochs,
|
||
normal_dataloader=normal_dataloader,
|
||
attack_dataloader=attack_dataloader,
|
||
model=model, evaluations=evaluations,
|
||
device=device, lr=learning_rate,
|
||
model_path=pth_name, train=train)
|
||
# 保存最佳检测结果的标签为csv文件
|
||
best_detection = pandas.DataFrame(best_detection)
|
||
best_detection.to_csv(f"./records/{self.start_time}/detection_result/{model_name}_{dataset_name}.csv", index=False)
|
||
for evaluation_name in evaluations:
|
||
if evaluation_name not in all_score:
|
||
all_score[evaluation_name] = []
|
||
all_score[evaluation_name].append(best_score[evaluation_name])
|
||
gc.collect()
|
||
logger.info(f"------------------------")
|
||
logger.info(f"{model_name} / {datasets[i]} 实验完毕")
|
||
for evaluation_name in all_score:
|
||
logger.info(f"{evaluation_name}: {'{:.3f}'.format(sum(all_score[evaluation_name]) / len(all_score[evaluation_name]))}")
|
||
logger.info(f"------------------------")
|
||
logger.info(f"------------华丽丽的分界线：{model_name} 实验结束------------")
|
||
except RuntimeError:
|
||
logger.error(traceback.format_exc())
|
||
return
|
||
|
||
|
||
if __name__ == '__main__':
|
||
pidpath = "/tmp/command_detection.pid"
|
||
app = Main(pidpath)
|
||
|