From 1b05affad60d360eef4aa22623aea5b1a5352db5 Mon Sep 17 00:00:00 2001 From: QueensGambit Date: Mon, 5 Aug 2024 18:03:28 +0200 Subject: [PATCH] Added processes commad-line param --- .../preprocessing/puzzles/puzzle_to_planes_converter.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/DeepCrazyhouse/src/preprocessing/puzzles/puzzle_to_planes_converter.py b/DeepCrazyhouse/src/preprocessing/puzzles/puzzle_to_planes_converter.py index 5feeb8ae..3fd61117 100644 --- a/DeepCrazyhouse/src/preprocessing/puzzles/puzzle_to_planes_converter.py +++ b/DeepCrazyhouse/src/preprocessing/puzzles/puzzle_to_planes_converter.py @@ -78,20 +78,21 @@ def sort_concat_data(data_dic: dict): return np.concatenate(data, axis=0) -def process_chunk(chunk_id: int, chunksize: int, df_chunk: pd.DataFrame, export_dir: Path): +def process_chunk(chunk_id: int, chunksize: int, df_chunk: pd.DataFrame, export_dir: Path, processes: int): """ Processes a data frame chunk by exporting all chess puzzle positions in this chunk. :param chunk_id: Unique id of the data chunk :param chunksize: Size of each chunk :param df_chunk: Data frame chunk :param export_dir: Export directory where the .zip files will be stored + :param processes: Number of processes return: None """ # engine = chess.engine.SimpleEngine.popen_uci(r"stockfish") logging.info("starting conversion to planes...") - pool = Pool() + pool = Pool(processes=processes) x_dic = {} y_value_dic = {} y_policy_dic = {} @@ -168,6 +169,7 @@ def _export_data(chunk_id, export_dir, phase_vector_dic, plys_to_end_dic, x_dic, parser.add_argument('--puzzle-csv-dir', type=str, default='./', help='Directory where the puzzle csv file is stored.') parser.add_argument('--export-dir', type=str, default='./', help='Directory where the .zip files will be exported to.') + parser.add_argument('--processes', type=int, default='4', help='Number of parallel processes.') args = parser.parse_args() @@ -199,4 +201,4 @@ def _export_data(chunk_id, export_dir, phase_vector_dic, plys_to_end_dic, x_dic, with pd.read_csv(puzzle_file_path, chunksize=chunksize) as reader: for chunk_id, df_chunk in enumerate(reader): print('chunk:', df_chunk) - process_chunk(chunk_id, chunksize, df_chunk, export_dir) + process_chunk(chunk_id, chunksize, df_chunk, export_dir, args.processes)