(input_path, output_path, cuda_list)
| 3 | import pandas as pd |
| 4 | |
| 5 | def split(input_path, output_path, cuda_list): |
| 6 | if type(cuda_list) == int: |
| 7 | cuda_list = [cuda_list] |
| 8 | df = pd.read_csv(input_path) |
| 9 | # df = df.sample(frac=1).reset_index(drop=True) |
| 10 | if not os.path.exists(output_path): |
| 11 | os.makedirs(output_path) |
| 12 | df_len = len(df) |
| 13 | cuda_list = list(cuda_list) |
| 14 | cuda_num = len(cuda_list) |
| 15 | for i in range(cuda_num): |
| 16 | start = i * df_len // cuda_num |
| 17 | end = (i+1) * df_len // cuda_num |
| 18 | df[start:end].to_csv(f'{output_path}/{cuda_list[i]}.csv', index=True) |
| 19 | |
| 20 | if __name__ == '__main__': |
| 21 | fire.Fire(split) |
nothing calls this directly
no outgoing calls
no test coverage detected