import os import numpy import torch.nn as nn import torch import torch.utils.data as tud import preprocess import evaluation from loguru import logger import time from tqdm import tqdm def create_dataloader(dataset_name: str, input_size: int = 1, output_size: int = 1, step: int = 1, batch_size: int = 1, time_index: bool = True, del_column_name: bool = True, preprocess_name: str = "standardization") -> (tud.DataLoader, tud.DataLoader): """ 针对一个数据集，构建Dataloader :param dataset_name: 数据集名称 :param input_size: 输入数据长度 :param output_size: 输出数据长度 :param step: 截取数据的窗口移动间隔 :param batch_size: batch的大小 :param time_index: True为第一列是时间戳，False为不 :param del_column_name: 文件中第一行为列名时，使用True :param preprocess_name: 预处理方法 :return: 训练数据与测试数据 """ ds = eval(f"preprocess.{preprocess_name}.MyDataset")(name=dataset_name.replace("/", "-"), train_path=f"./dataset/{dataset_name}/train.csv", test_path=f"./dataset/{dataset_name}/test.csv", input_size=input_size, output_size=output_size, step=step, time_index=time_index, del_column_name=del_column_name) normal_dl = tud.DataLoader(dataset=ds, batch_size=batch_size, shuffle=True) ds.mode = "test" attack_dl = tud.DataLoader(dataset=ds, batch_size=batch_size, shuffle=False) return normal_dl, attack_dl def create_all_dataloader(datasets: [str], input_size: int = 1, output_size: int = 1, step: int = 1, batch_size: int = 1, time_index: bool = True, del_column_name: bool = True, preprocess_name: str = "standardization") -> [{}]: """ 对所有数据集构建dataloader :param datasets: 数据集列表 :param input_size: 输入数据长度 :param output_size: 输出数据长度 :param step: 截取数据的窗口移动间隔 :param batch_size: batch的大小 :param time_index: True为第一列是时间戳，False为不。 :param del_column_name: 文件中第一行为列名时，使用True :param preprocess_name: 预处理方法 :return: 所有数据集的dataloader构建结果 """ all_dataloader = [] for dataset_name in datasets: logger.info(f'开始建立 dataloader {dataset_name}') if "train.csv" in os.listdir(f"./dataset/{dataset_name}"): normal_dl, attack_dl = create_dataloader(dataset_name=dataset_name, input_size=input_size, output_size=output_size, step=step, batch_size=batch_size, time_index=time_index, del_column_name=del_column_name, preprocess_name=preprocess_name) all_dataloader.append([{ 'dataset_name': dataset_name, 'normal': normal_dl, 'attack': attack_dl }]) else: all_sub_dataloader = [] for sub_dataset_dir in os.listdir(f"./dataset/{dataset_name}"): sub_dataset_name = f"{dataset_name}/{sub_dataset_dir}" normal_dl, attack_dl = create_dataloader(dataset_name=sub_dataset_name, input_size=input_size, output_size=output_size, step=step, batch_size=batch_size, del_time=del_time, del_column_name=del_column_name) all_sub_dataloader.append({ 'dataset_name': sub_dataset_name.replace("/", "-"), 'normal': normal_dl, 'attack': attack_dl }) all_dataloader.append(all_sub_dataloader) logger.info(f'完成建立 dataloader {dataset_name}') return all_dataloader class EvaluationScore: def __init__(self, evaluations: [str], attack=1): """ 用于自动划分阈值并进行批量评估 :param evaluations: 使用的评估方法名称（需在evaluation文件夹中进行定义） :param attack: 异常的标签，0 or 1 """ self.time = 0 self.f1 = 0 self.f1_pa = 0 self.f_tad = 0 self.attack = attack self.normal = 1 - attack self._total_y_loss = None self._total_label = None self._total_pred_label = None self.true_pred_df = None self.true_pred_dict = None self.evaluations = evaluations self.scores = {} def add_data(self, y_loss, true_label, pred_label=None): """ 添加每个batch的数据 :param y_loss: 数据偏差 :param true_label: 真实数据标签 :param pred_label: 预测标签 """ if pred_label is not None: if self._total_label is None and self._total_pred_label is None: self._total_label = true_label self._total_pred_label = pred_label else: self._total_label = torch.cat([self._total_label, true_label], dim=0) self._total_pred_label = torch.cat([self._total_pred_label, pred_label], dim=0) return y_loss = y_loss.view(-1).cpu().detach().numpy() true_label = true_label.view(-1).cpu().detach().numpy() if self._total_y_loss is None and self._total_label is None: self._total_y_loss = y_loss self._total_label = true_label return self._total_y_loss = numpy.concatenate((self._total_y_loss, y_loss), axis=0) self._total_label = numpy.concatenate((self._total_label, true_label), axis=0) def best_threshold(self, true_label: list, y_loss: list) -> dict: ret = {} for func_name in self.evaluations: threshold_max = max(y_loss) threshold_min = 0 best_threshold = 0 for _ in range(5): threshold_list = [threshold_max - i * (threshold_max - threshold_min) / 10 for i in range(11)] f1_list = [] for threshold_one in threshold_list: prediction_loss = numpy.where(numpy.array(y_loss) > threshold_one, self.attack, self.normal) f1 = eval(f"evaluation.{func_name}.evaluate")(y_true=true_label, y_pred=prediction_loss.tolist()) f1_list.append(f1) ind = f1_list.index(max(f1_list)) best_threshold = threshold_list[ind] if ind == 0: threshold_max = threshold_list[ind] threshold_min = threshold_list[ind+1] elif ind == len(threshold_list)-1: threshold_max = threshold_list[ind-1] threshold_min = threshold_list[ind] else: threshold_max = threshold_list[ind-1] threshold_min = threshold_list[ind+1] ret[func_name] = best_threshold return ret def auto_threshold(self): if self._total_pred_label is not None: return self._total_y_loss[numpy.isnan(self._total_y_loss)] = 0 self._total_y_loss = self._total_y_loss / max(self._total_y_loss) thresholds = self.best_threshold( self._total_label.reshape(-1).data.tolist(), self._total_y_loss.reshape(-1).data.tolist()) self.true_pred_dict = { 'true': self._total_label.squeeze().tolist() } for func_name in thresholds: self.true_pred_dict[func_name] = \ numpy.where(self._total_y_loss > thresholds[func_name], self.attack, self.normal).squeeze().tolist() # self.true_pred_df = pandas.DataFrame(self.true_pred_dict) for func_name in self.true_pred_dict: if func_name == "true": continue self.scores[func_name] = self.get_score(func_name) def get_score(self, func_name): if self._total_pred_label is not None: return eval(f"evaluation.{func_name}.evaluate")(self._total_label.reshape(-1).tolist(), self._total_pred_label.reshape(-1).tolist()) return eval(f"evaluation.{func_name}.evaluate")(self._total_label.reshape(-1).tolist(), self.true_pred_dict[f"{func_name}"]) def __str__(self): res = "" for func_name in self.scores: res += f"{func_name}={self.scores[func_name]:.3f} " return res[:-1] def train_model(epoch: int, optimizer: torch.optim, dataloader: tud.DataLoader, model: nn.Module, device: str = "cpu") -> (nn.Module, str): """ 训练模型 :param epoch: 当前训练轮数 :param optimizer: 优化器 :param dataloader: 数据集 :param model: 模型 :param device: 训练设备使用cpu还是gpu :return: 训练完成的模型；训练完成需要输出的信息 """ model.train() avg_loss = [] dataloader.dataset.mode = "train" start_time = time.time() with tqdm(total=len(dataloader), ncols=150) as _tqdm: _tqdm.set_description(f'（进度条部分不会写进本地日志）epoch:{epoch}，训练进度') for data in dataloader: x = data[0].to(device) y_true = data[2].to(device) optimizer.zero_grad() loss = model.loss(x=x, y_true=y_true, epoch=epoch, device=device) avg_loss.append(loss) optimizer.step() _tqdm.set_postfix(loss='{:.6f}'.format(sum(avg_loss) / len(avg_loss))) _tqdm.update(1) end_time = time.time() info = f"epoch={epoch}, average loss={'{:.6f}'.format(sum(avg_loss) / len(avg_loss))}, " \ f"train time={'{:.1f}'.format(end_time-start_time)}s" return model, info def test_model(dataloader: tud.DataLoader, model: nn.Module, evaluations: [str], device: str = "cpu") -> \ (EvaluationScore, str): """ 测试模型 :param dataloader: 数据集 :param model: 模型 :param device: 训练设备使用cpu还是gpu :return: 评估分数；测试完成需要输出的信息 """ es = EvaluationScore(evaluations) model.eval() dataloader.dataset.mode = "test" start_time = time.time() with tqdm(total=len(dataloader), ncols=150) as _tqdm: _tqdm.set_description(f'（进度条部分不会写进本地日志）测试进度') with torch.no_grad(): for data in dataloader: x = data[0].to(device) y_true = data[2].to(device) label_true = data[1].int().to(device) y_loss, label_pred = model.detection(x=x, y_true=y_true, device=device) if label_pred is not None: es.add_data(y_loss=None, true_label=label_true, pred_label=label_pred) else: es.add_data(y_loss=y_loss, true_label=label_true, pred_label=None) _tqdm.update(1) end_time = time.time() es.auto_threshold() es_score = es.__str__().replace(" ", ", ") info = f"{es_score}, test time={'{:.1f}'.format(end_time-start_time)}s" return es, info def train_and_test_model(start_time: str, epochs: int, normal_dataloader: tud.DataLoader, attack_dataloader: tud.DataLoader, model: nn.Module, evaluations: [str], device: str = "cpu", lr: float = 1e-4, model_path: str = None, train: bool = True) -> (dict, dict): """ 训练与测试 :param start_time: 实验的开始时间。此处用于寻找存放路径。 :param epochs: 总共训练轮数 :param normal_dataloader: 训练数据集 :param attack_dataloader: 测试数据集 :param model: 模型 :param evaluations: 评估方法 :param device: 设备 :param lr: 学习率 :param model_path: 模型参数文件路径 :param train: 是否训练，如果为否，则仅进行测试 :return: 各个评估方法的最佳分数、各个评估方法最佳情况下的检测标签 """ dataset_name = normal_dataloader.dataset.name optimizer = torch.optim.Adam(model.parameters(), lr=lr) if model_path: try: checkpoint = torch.load(model_path) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) logger.info(f"模型参数文件{model_path}加载成功") except: logger.warning(f"模型参数文件{model_path}加载失败，将训练新模型") logger.info(f"模型：{model.name}，数据集：{dataset_name}，设备：{device}，训练开始") best_score = {} best_detection = {} if train: logger.info(f"模式：训练并测试") for epoch in range(1, epochs+1): model, train_info = train_model(epoch=epoch, optimizer=optimizer, dataloader=normal_dataloader, model=model, device=device) es, test_info = test_model(dataloader=attack_dataloader, model=model, evaluations=evaluations, device=device) logger.info(f"{train_info}, {test_info}") es_score = es.__str__().replace(" ", "_") torch.save({ 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict() }, f'./records/{start_time}/model/model={model.name}_dataset={dataset_name}_epoch={epoch}_{es_score}.pth') for func_name in es.scores: if func_name not in best_score or es.scores[func_name] > best_score[func_name]: best_score[func_name] = es.scores[func_name] best_detection[func_name] = es.true_pred_dict[func_name] best_detection["true"] = es.true_pred_dict["true"] else: logger.info(f"模式：仅进行测试") es, test_info = test_model(dataloader=attack_dataloader, model=model, evaluations=evaluations, device=device) logger.info(test_info) for func_name in es.scores: best_score[func_name] = es.scores[func_name] best_detection[func_name] = es.true_pred_dict[func_name] best_detection["true"] = es.true_pred_dict["true"] return best_score, best_detection