| 57 | |
| 58 | # Data partition according to the rank |
| 59 | def partition(global_rank, world_size, train_x, train_y, val_x, val_y): |
| 60 | # Partition training data |
| 61 | data_per_rank = train_x.shape[0] // world_size |
| 62 | idx_start = global_rank * data_per_rank |
| 63 | idx_end = (global_rank + 1) * data_per_rank |
| 64 | train_x = train_x[idx_start:idx_end] |
| 65 | train_y = train_y[idx_start:idx_end] |
| 66 | |
| 67 | # Partition evaluation data |
| 68 | data_per_rank = val_x.shape[0] // world_size |
| 69 | idx_start = global_rank * data_per_rank |
| 70 | idx_end = (global_rank + 1) * data_per_rank |
| 71 | val_x = val_x[idx_start:idx_end] |
| 72 | val_y = val_y[idx_start:idx_end] |
| 73 | return train_x, train_y, val_x, val_y |
| 74 | |
| 75 | |
| 76 | # Function to all reduce NUMPY accuracy and loss from multiple devices |