| 74 | |
| 75 | class ReshapeWithInput(Pipeline): |
| 76 | def __init__(self, device, batch_size, use_wildcard, num_threads=3, device_id=0, num_gpus=1): |
| 77 | super(ReshapeWithInput, self).__init__( |
| 78 | batch_size, num_threads, device_id, seed=7865, exec_async=False, exec_pipelined=False |
| 79 | ) |
| 80 | self.device = device |
| 81 | self.input = ops.readers.Caffe( |
| 82 | path=caffe_db_folder, shard_id=device_id, num_shards=num_gpus |
| 83 | ) |
| 84 | self.decode = ops.decoders.Image(device="cpu", output_type=types.RGB) |
| 85 | fn = CollapseChannelsWildcard if use_wildcard else CollapseChannels |
| 86 | self.gen_shapes = ops.PythonFunction(function=fn) |
| 87 | self.reshape = ops.Reshape(device=device, layout="ab") |
| 88 | |
| 89 | def define_graph(self): |
| 90 | jpegs, labels = self.input(name="Reader") |