58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
|
import configparser
|
|||
|
|
|
|||
|
|
|
|||
|
|
def read_and_group_data_by_class(input_file):
|
|||
|
|
data_by_class = {}
|
|||
|
|
with open(input_file + '_attack.txt', 'r') as file:
|
|||
|
|
for line in file:
|
|||
|
|
parts = line.strip().split()
|
|||
|
|
cls = parts[0]
|
|||
|
|
if cls not in data_by_class:
|
|||
|
|
data_by_class[cls] = []
|
|||
|
|
data_by_class[cls].append(parts)
|
|||
|
|
return data_by_class
|
|||
|
|
|
|||
|
|
def split_data(data, n):
|
|||
|
|
# 确保即使数据不能均匀分配也能处理
|
|||
|
|
total_length = len(data)
|
|||
|
|
split_length = max(total_length // n, 1)
|
|||
|
|
remainder = total_length % n
|
|||
|
|
splits = []
|
|||
|
|
start_idx = 0
|
|||
|
|
for i in range(n):
|
|||
|
|
if i < remainder:
|
|||
|
|
end_idx = start_idx + split_length + 1
|
|||
|
|
else:
|
|||
|
|
end_idx = start_idx + split_length
|
|||
|
|
splits.append(data[start_idx:end_idx])
|
|||
|
|
start_idx = end_idx
|
|||
|
|
return splits
|
|||
|
|
|
|||
|
|
def write_split_data(split_data_by_class, n, data_path):
|
|||
|
|
for i in range(n):
|
|||
|
|
output_file = f'{data_path}_attack{i}.txt'
|
|||
|
|
with open(output_file, 'w') as f:
|
|||
|
|
for cls, data_splits in split_data_by_class.items():
|
|||
|
|
if i < len(data_splits): # 检查索引以防止越界
|
|||
|
|
for line in data_splits[i]:
|
|||
|
|
f.write(" ".join(line) + "\n")
|
|||
|
|
|
|||
|
|
def split_and_save_data(data_path, n):
|
|||
|
|
# Step 1: Read and group data by class
|
|||
|
|
data_by_class = read_and_group_data_by_class(data_path)
|
|||
|
|
|
|||
|
|
# Step 2: Split data for each class
|
|||
|
|
split_data_by_class = {cls: split_data(data, n) for cls, data in data_by_class.items()}
|
|||
|
|
|
|||
|
|
# Step 3: Write split data into files
|
|||
|
|
write_split_data(split_data_by_class, n, data_path)
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# 临时添加配置文件类别种类
|
|||
|
|
config = configparser.ConfigParser()
|
|||
|
|
config.read("run_parallel_config.ini")
|
|||
|
|
file_name = config['MODEL']['run_tag']
|
|||
|
|
data_path = 'data/' + file_name + '/' + file_name
|
|||
|
|
n = 5 # 假设我们想将数据均分到7个文件中，最多7个
|
|||
|
|
split_and_save_data(data_path, n)
|