Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

load_npy_from_path: improve load balancing #1602

Merged
merged 1 commit into from
Aug 9, 2024
Merged

load_npy_from_path: improve load balancing #1602

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions heat/core/io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1183,11 +1183,14 @@ def load_npy_from_path(
raise RuntimeError("Number of processes can't exceed number of files")

rank = MPI_WORLD.rank
n_for_procs = n_files // process_number
idx = rank * n_for_procs
if rank + 1 == process_number:
n_for_procs += n_files % process_number
if rank < (n_files % process_number):
n_for_procs = n_files // process_number + 1
idx = rank * n_for_procs
else:
n_for_procs = n_files // process_number
idx = rank * n_for_procs + (n_files % process_number)
array_list = [np.load(path + "/" + element) for element in file_list[idx : idx + n_for_procs]]

larray = np.concatenate(array_list, split)
larray = torch.from_numpy(larray)

Expand Down
Loading