Dali: Stuck when running ExternalSourcePipeline.run()

[zhengrchan]

Hi,
I'm using ExternalSourcePipeline to load 3755359images/epoch. When I run

eii = ExternalInputIterator(root, image_dir, batch_size)
iterator = iter(eii)
pip_train = ExternalSourcePipeline(iterator, batch_size, num_threads=num_threads, device_id=local_rank, input_shape=input_shape)
pip_train.build()
print('pip_train building, len of epoch: %d' % iterator.size)
pipe_out = pip_train.run()

the process stuck at pipe_out = pip_train.run(), and I have to use 'kill PID` to terminate the process.

And the same thing happens when I run train_loader = DALIClassificationIterator(pip_train, size=iterator.size).

My config:

• pytorch: 1.1.0 with cuda9.0 and cudnn7.5.1
• cudatoolkit: 9.0
• and I use pip install --extra-index-url https://developer.download.nvidia.com/compute/redist/cuda/9.0 nvidia-dali to install Nvidia-DALI

Here is my code

class ExternalInputIterator(object):
    def __init__(self, root, images_dir, batch_size):
        self.batch_size = batch_size

        with open(os.path.join(images_dir), 'r') as fd:
            imgs = fd.readlines()

        imgs = [os.path.join(root, img[:-1]) for img in imgs]
        self.imgs = np.random.permutation(imgs)
        self.len = len(self.imgs)

    def __iter__(self):
        self.i = 0
        self.n = len(self.imgs)
        return self

    def __next__(self):
        batch = []
        labels = []
        for _ in range(self.batch_size):
            sample = self.imgs[self.i].split()
            with open(sample[0], 'rb') as f:
                batch.append(np.frombuffer(f.read(), dtype=np.uint8))
                labels.append(np.array([sample[1]], dtype=np.uint8))
            self.i = (self.i + 1) % self.n
        return (batch, labels)  

    def __len__(self):
        return self.len

    @property
    def size(self,):
        return self.len

    next = __next__


class ExternalSourcePipeline(Pipeline):
    def __init__(self, iterator, batch_size, num_threads, device_id, input_shape=(3,248,248)):
        super(ExternalSourcePipeline, self).__init__(batch_size, num_threads, device_id, seed=12)
        self.iterator = iterator
        self.input = ops.ExternalSource()
        self.input_label = ops.ExternalSource()
        self.decode = ops.ImageDecoder(device='mixed', output_type=types.RGB)
        self.cmnp = ops.CropMirrorNormalize(device="gpu",
                                            output_dtype=types.FLOAT,
                                            output_layout=types.NCHW,
                                            crop = (input_shape[1], input_shape[2]),
                                            image_type=types.RGB,
                                            mean=[0.485 * 255., 0.456 * 255., 0.406 * 255.],
                                            std=[0.229 * 255., 0.224 * 255., 0.225 * 255.]
                                            )
        self.coin = ops.CoinFlip(probability=0.5)

    def define_graph(self):
        rng = self.coin()
        self.images = self.input()
        self.labels = self.input_label()
        images = self.decode(self.images)
        output = self.cmnp(images.gpu(), mirror=rng)
        return (output, self.labels)

    def iter_step(self):
        (images, labels) = self.iterator.next()
        self.feed_input(self.images, images)
        self.feed_input(self.labels, labels)

[JanuszL]

iter_step -> iter_setup

[zhengrchan]

iter_step -> iter_setup

Thank you so much!