120 lines
5.9 KiB
Python
120 lines
5.9 KiB
Python
from torch.utils.data import Dataset
|
||
from torch import float32, Tensor
|
||
from numpy import array, where
|
||
|
||
|
||
class MyDataset(Dataset):
|
||
def __init__(self, name: str, train_path: str = None, test_path: str = None, input_size: int = 1,
|
||
output_size: int = 1, step: int = 1, mode: str = 'train', time_index: bool = True,
|
||
del_column_name: bool = True):
|
||
"""
|
||
可以将csv文件批量转成tensor
|
||
:param name: 数据集名称。
|
||
:param train_path: 训练数据集路径。
|
||
:param test_path: 测试数据集路径。
|
||
:param input_size: 输入数据长度。
|
||
:param output_size: 输出数据长度。
|
||
:param step: 截取数据的窗口移动间隔。
|
||
:param mode: train或者test，用于指示使用训练集数据还是测试集数据。
|
||
:param time_index: True为第一列是时间戳，False为不。
|
||
:param del_column_name: 文件中第一行为列名时，使用True。
|
||
"""
|
||
self.name = name
|
||
self.input_size = input_size
|
||
self.output_size = output_size
|
||
self.del_column_name = del_column_name
|
||
self.step = step
|
||
self.mode = mode
|
||
self.time_index = time_index
|
||
self.train_inputs, self.train_labels, self.train_outputs, self.test_inputs, self.test_labels, self.test_outputs\
|
||
= self.parse_data(train_path, test_path)
|
||
self.train_inputs = Tensor(self.train_inputs).to(float32) if self.train_inputs is not None else None
|
||
self.train_labels = Tensor(self.train_labels).to(float32) if self.train_labels is not None else None
|
||
self.train_outputs = Tensor(self.train_outputs).to(float32) if self.train_outputs is not None else None
|
||
self.test_inputs = Tensor(self.test_inputs).to(float32) if self.test_inputs is not None else None
|
||
self.test_labels = Tensor(self.test_labels).to(float32) if self.test_labels is not None else None
|
||
self.test_outputs = Tensor(self.test_outputs).to(float32) if self.test_outputs is not None else None
|
||
|
||
def parse_data(self, train_path: str = None, test_path: str = None):
|
||
if train_path is None and test_path is None:
|
||
raise ValueError("train_path is None and test_path is None.")
|
||
|
||
mean = None
|
||
deviation = None
|
||
train_data_input, train_label, train_data_output = None, None, None
|
||
test_data_input, test_label, test_data_output = None, None, None
|
||
|
||
# 读取训练集数据
|
||
if train_path:
|
||
train_data = []
|
||
train_label = []
|
||
with open(train_path, 'r', encoding='utf8') as f:
|
||
if self.del_column_name is True:
|
||
data = f.readlines()[1:]
|
||
else:
|
||
data = f.readlines()
|
||
train_data.extend([list(map(float, line.strip().split(','))) for line in data])
|
||
train_label.extend([0 for _ in data])
|
||
train_np = array(train_data)
|
||
if self.time_index:
|
||
train_np[:, 0] = train_np[:, 0] % 86400
|
||
mean = train_np.mean(axis=0) # 计算平均数
|
||
deviation = train_np.std(axis=0) # 计算标准差
|
||
deviation = where(deviation != 0, deviation, 1)
|
||
train_np = (train_np - mean) / deviation # 标准化
|
||
train_data = train_np.tolist()
|
||
train_data_input, train_data_output, train_label = self.cut_data(train_data, train_label)
|
||
|
||
# 读取测试集数据
|
||
if test_path:
|
||
test_data = []
|
||
test_label = []
|
||
with open(test_path, 'r', encoding='utf8') as f:
|
||
if self.del_column_name is True:
|
||
data = f.readlines()[1:]
|
||
else:
|
||
data = f.readlines()
|
||
test_data.extend([list(map(float, line.strip().split(',')))[:-1] for line in data])
|
||
test_label.extend([int(line.strip().split(',')[-1]) for line in data])
|
||
test_np = array(test_data)
|
||
if self.time_index:
|
||
test_np[:, 0] = test_np[:, 0] % 86400
|
||
# mean = test_np.mean(axis=0) # 计算平均数
|
||
# deviation = test_np.std(axis=0) # 计算标准差
|
||
# deviation = where(deviation != 0, deviation, 1)
|
||
test_np = (test_np - mean) / deviation # 标准化
|
||
test_data = test_np.tolist()
|
||
# 自动判断是否需要反转标签。异常标签统一认为是1，当异常标签超过一半时，需反转标签
|
||
if sum(test_label) > 0.5*len(test_label):
|
||
test_label = (1-array(test_label)).tolist()
|
||
test_data_input, test_data_output, test_label = self.cut_data(test_data, test_label)
|
||
|
||
return train_data_input, train_label, train_data_output, test_data_input, test_label, test_data_output
|
||
|
||
def cut_data(self, data: [[float]], label: [int]):
|
||
n = 0
|
||
input_data, output_data, anomaly_label = [], [], []
|
||
while n + self.input_size + self.output_size <= len(data):
|
||
input_data.append(data[n: n + self.input_size])
|
||
output_data.append(data[n + self.input_size: n + self.input_size + self.output_size])
|
||
anomaly_label.append([max(label[n + self.input_size: n + self.input_size + self.output_size])])
|
||
n = n + self.step
|
||
return input_data.copy(), output_data.copy(), anomaly_label.copy()
|
||
|
||
def __len__(self):
|
||
if self.mode == 'train':
|
||
return self.train_inputs.shape[0]
|
||
elif self.mode == 'test':
|
||
return self.test_inputs.shape[0]
|
||
|
||
def __getitem__(self, idx):
|
||
if self.mode == 'train':
|
||
return self.train_inputs[idx], self.train_labels[idx], self.train_outputs[idx]
|
||
elif self.mode == 'test':
|
||
return self.test_inputs[idx], self.test_labels[idx], self.test_outputs[idx]
|
||
|
||
|
||
if __name__ == "__main__":
|
||
app = MyDataset('../dataset/SWAT/train.csv', test_path='../dataset/SWAT/test.csv', input_size=3)
|
||
print(app)
|