Create dataset Inputs: batch_size: batch size data_path: path of your MindRecord files device_num: total device number rank: current rank id drop: whether drop remainder eod_reset: whether enable position reset and attention mask reset
(batch_size, data_path, args_opt, device_num=1, rank=0, drop=True, full_batch=False,
data_start_index=0,
eod_reset=False, eod_id=50256, column_name='input_ids', epoch=1, num_samples=None,
train_and_eval=False, val_ratio=0)
| 78 | |
| 79 | |
| 80 | def create_dataset(batch_size, data_path, args_opt, device_num=1, rank=0, drop=True, full_batch=False, |
| 81 | data_start_index=0, |
| 82 | eod_reset=False, eod_id=50256, column_name='input_ids', epoch=1, num_samples=None, |
| 83 | train_and_eval=False, val_ratio=0): |
| 84 | """ |
| 85 | Create dataset |
| 86 | Inputs: |
| 87 | batch_size: batch size |
| 88 | data_path: path of your MindRecord files |
| 89 | device_num: total device number |
| 90 | rank: current rank id |
| 91 | drop: whether drop remainder |
| 92 | eod_reset: whether enable position reset and attention mask reset |
| 93 | eod_id: the id for <EOD> |
| 94 | column_name: the column name of the mindrecord file. Default is input_ids |
| 95 | epoch: The repeat times of the dataset |
| 96 | Returns: |
| 97 | dataset_restore: the dataset for training or evaluating |
| 98 | """ |
| 99 | ds.config.set_seed(1) |
| 100 | # Control the size of data queue in the consideration of the memory |
| 101 | ds.config.set_prefetch_size(1) |
| 102 | |
| 103 | if full_batch: |
| 104 | # no need to slice from the inputs |
| 105 | rank = 0 |
| 106 | dis = batch_size |
| 107 | else: |
| 108 | # Each card slice a small batch from the full batch |
| 109 | dis = int(batch_size / device_num) |
| 110 | if batch_size % device_num != 0: |
| 111 | raise ValueError( |
| 112 | f"batch size {batch_size} should be a multiple of device number {device_num}." |
| 113 | " You should change the args: per_batch_size." |
| 114 | ) |
| 115 | |
| 116 | skip_num = args_opt.has_trained_steps * dis |
| 117 | # skip_num = 0 |
| 118 | num_parallel_workers = 4 |
| 119 | train_data = get_code_data_train(data_path, args_opt, skip_num=(skip_num // num_parallel_workers)) |
| 120 | if train_and_eval: |
| 121 | val_data = get_code_data_eval("/home/work/sfs/xx/data_valid", |
| 122 | args_opt) # TODO: set as current validation set path |
| 123 | else: |
| 124 | val_data = None |
| 125 | |
| 126 | dataset_train = ds.GeneratorDataset(train_data, column_names=[column_name], num_samples=num_samples, |
| 127 | num_shards=device_num, shard_id=rank, shuffle=True, |
| 128 | num_parallel_workers=num_parallel_workers) |
| 129 | if train_and_eval: |
| 130 | dataset_val = ds.GeneratorDataset(val_data, column_names=[column_name], num_samples=num_samples, |
| 131 | num_shards=device_num, shard_id=rank, shuffle=True, |
| 132 | num_parallel_workers=num_parallel_workers) |
| 133 | else: |
| 134 | dataset_val = None |
| 135 | type_cast_op = C.TypeCast(mstype.int32) |
| 136 | type_cast_op_float = C.TypeCast(mstype.float16) |
| 137 |
no test coverage detected