| 153 | |
| 154 | |
| 155 | def to_image_dataset(image_pipeline_desc, device_str): |
| 156 | dataset_pipeline, shapes, dtypes = image_pipeline_desc |
| 157 | with tf.device(device_str): |
| 158 | dali_dataset = dali_tf.DALIDataset( |
| 159 | pipeline=dataset_pipeline, |
| 160 | batch_size=dataset_pipeline.batch_size, |
| 161 | output_shapes=shapes, |
| 162 | output_dtypes=dtypes, |
| 163 | num_threads=dataset_pipeline.num_threads, |
| 164 | device_id=dataset_pipeline.device_id, |
| 165 | ) |
| 166 | return dali_dataset |
| 167 | |
| 168 | |
| 169 | def get_dali_dataset_from_pipeline(pipeline_desc, device, device_id, to_dataset=to_image_dataset): |