(global_rank,
world_size,
local_rank,
max_epoch,
batch_size,
model,
data,
dir_path,
sgd,
graph,
verbosity,
dist_option='plain',
spars=None,
peft_type='None',
precision='float32',
)
| 95 | return X |
| 96 | |
| 97 | def run(global_rank, |
| 98 | world_size, |
| 99 | local_rank, |
| 100 | max_epoch, |
| 101 | batch_size, |
| 102 | model, |
| 103 | data, |
| 104 | dir_path, |
| 105 | sgd, |
| 106 | graph, |
| 107 | verbosity, |
| 108 | dist_option='plain', |
| 109 | spars=None, |
| 110 | peft_type='None', |
| 111 | precision='float32', |
| 112 | ): |
| 113 | dev = device.get_default_device() |
| 114 | dev.SetRandSeed(0) |
| 115 | np.random.seed(0) |
| 116 | |
| 117 | if data == "mnist": |
| 118 | from examples.data import mnist |
| 119 | train_x, train_y, val_x, val_y = mnist.load(dir_path) |
| 120 | else: |
| 121 | raise ValueError(f"`r`Not support dataset {data}") |
| 122 | |
| 123 | |
| 124 | num_channels = train_x.shape[1] |
| 125 | image_size = train_x.shape[2] |
| 126 | data_size = np.prod(train_x.shape[1:train_x.ndim]).item() |
| 127 | num_classes = (np.max(train_y) + 1).item() |
| 128 | |
| 129 | if model == "mlp": |
| 130 | from examples.model import mlp |
| 131 | model = mlp.create_model(in_features=data_size, perceptron_size=16, num_classes=num_classes) |
| 132 | elif model == "cnn": |
| 133 | from examples.model import cnn |
| 134 | model = cnn.create_model(num_channels=num_channels, num_classes=num_classes) |
| 135 | else: |
| 136 | raise ValueError(f"`r`Not support model {model}") |
| 137 | |
| 138 | # For distributed training, sequential has better performance |
| 139 | if hasattr(sgd, "communicator"): |
| 140 | DIST = True |
| 141 | sequential = True |
| 142 | else: |
| 143 | DIST = False |
| 144 | sequential = False |
| 145 | |
| 146 | if DIST: |
| 147 | train_x, train_y, val_x, val_y = partition(global_rank, world_size, |
| 148 | train_x, train_y, val_x, |
| 149 | val_y) |
| 150 | |
| 151 | if model.dimension == 4: |
| 152 | tx = tensor.Tensor( |
| 153 | (batch_size, num_channels, model.input_size, model.input_size), dev, singa_dtype[precision]) |
| 154 | elif model.dimension == 2: |
no test coverage detected