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, epoch=1, num_samples=None, train_and_eval=False, val_ratio=0)
| 80 | |
| 81 | |
| 82 | def create_dataset(batch_size, data_path, args_opt, device_num=1, rank=0, drop=True, full_batch=False, |
| 83 | data_start_index=0, |
| 84 | eod_reset=False, eod_id=50256, epoch=1, num_samples=None, train_and_eval=False, val_ratio=0): |
| 85 | """ |
| 86 | Create dataset |
| 87 | Inputs: |
| 88 | batch_size: batch size |
| 89 | data_path: path of your MindRecord files |
| 90 | device_num: total device number |
| 91 | rank: current rank id |
| 92 | drop: whether drop remainder |
| 93 | eod_reset: whether enable position reset and attention mask reset |
| 94 | eod_id: the id for <EOD> |
| 95 | column_name: the column name of the mindrecord file. Default is input_ids |
| 96 | epoch: The repeat times of the dataset |
| 97 | Returns: |
| 98 | dataset_restore: the dataset for training or evaluating |
| 99 | """ |
| 100 | ds.config.set_seed(1) |
| 101 | # Control the size of data queue in the consideration of the memory |
| 102 | ds.config.set_prefetch_size(1) |
| 103 | |
| 104 | if full_batch: |
| 105 | # no need to slice from the inputs |
| 106 | rank = 0 |
| 107 | dis = batch_size |
| 108 | else: |
| 109 | # Each card slice a small batch from the full batch |
| 110 | dis = int(batch_size / device_num) |
| 111 | if batch_size % device_num != 0: |
| 112 | raise ValueError( |
| 113 | f"batch size {batch_size} should be a multiple of device number {device_num}." |
| 114 | " You should change the args: per_batch_size." |
| 115 | ) |
| 116 | |
| 117 | # skip_num = args_opt.has_trained_steps * dis |
| 118 | # skip_num = 0 |
| 119 | num_parallel_workers = 4 |
| 120 | train_data = get_code_data(data_path, 'train', args_opt) |
| 121 | if train_and_eval: |
| 122 | val_data = get_code_data(data_path, 'val', args_opt) |
| 123 | else: |
| 124 | val_data = None |
| 125 | |
| 126 | dataset_train = ds.GeneratorDataset(train_data, column_names=['input_ids', 'loss_mask'], 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=['input_ids', 'loss_mask'], 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 | type_cast_op_float2 = C.TypeCast(mstype.float32) |
| 138 | |
| 139 | map_func = ( |
no test coverage detected