(batch_size, pipe_type, label_type, num_gpus=1)
| 72 | |
| 73 | |
| 74 | def get_batch_dali(batch_size, pipe_type, label_type, num_gpus=1): |
| 75 | pipes = [ |
| 76 | pipe_type(batch_size=batch_size, num_threads=2, device_id=device_id, num_gpus=num_gpus) |
| 77 | for device_id in range(num_gpus) |
| 78 | ] |
| 79 | |
| 80 | daliop = dali_tf.DALIIterator() |
| 81 | images = [] |
| 82 | labels = [] |
| 83 | for d in range(NUM_GPUS): |
| 84 | with tf.device("/gpu:%i" % d): |
| 85 | image, label = daliop( |
| 86 | pipeline=pipes[d], |
| 87 | shapes=[(batch_size, 3, 227, 227), ()], |
| 88 | dtypes=[tf.int32, label_type], |
| 89 | device_id=d, |
| 90 | ) |
| 91 | images.append(image) |
| 92 | labels.append(label) |
| 93 | |
| 94 | return [images, labels] |
| 95 | |
| 96 | |
| 97 | def test_dali_tf_op(pipe_type=CaffeReadPipeline, batch_size=16, iterations=32): |
no test coverage detected