How to create DataPipe that best optimize the map-transform function that supports batching?

alvations: <p>Given a transformation function that <strong>I can’t change</strong>, e.g. <code>autobot_vectorize</code>. It</p>
<ul>
<li>takes in a list of N inputs and</li>
<li>output a tensor of N x 3 dimensions.</li>
</ul>
<pre><code class="lang-auto">def autobot_vectorize(imgfiles):
    # This vectorizer takes N imgfiles and return the 
    # a tensor of size N x 3.
    # For example, we are simulating this with torch.rand(N, 3)
    """
    &gt;&gt;&gt; autobot_vectorize(['train/0/47215.png', 'train/0/56789.png'])
    tensor([[1.23, 4.56, 7.89], [9.99, 8.88, 7.77]])
    """
    return torch.rand(len(imgfiles), 3)
</code></pre>
<p>And given a dataset like:</p>
<pre><code class="lang-auto">content = """imagefile,imagefile2,label
train/0/16585.png,0
train/0/56789.png,0
train/0/47215.png,1
train/0/93155.png,1"""

with open('myfile.csv', 'w') as fout:
    fout.write(content)
</code></pre>
<p>The objective is to create a Pytorch datapipe that will output something like this:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; dp = MyDataPipe(csvfile)
&gt;&gt;&gt; for row in dp.train_dataloader:
...     print(row)
...
(tensor([1.23, 4.56, 7.89]), 0)
(tensor([9.87, 6.54, 3.21]), 1)
(tensor([9.99, 8.88, 7.77]), 0)
(tensor([1.11, 2.22, 9.87]), 1)
</code></pre>
<h2>
<a name="what-is-the-a-simple-way-to-create-the-dataloader-from-the-datapipe-that-best-optimize-the-batching-function-of-the-autobot_vectorize-function-1" class="anchor" href="#what-is-the-a-simple-way-to-create-the-dataloader-from-the-datapipe-that-best-optimize-the-batching-function-of-the-autobot_vectorize-function-1"></a>What is the a simple way to create the DataLoader from the DataPipe that best optimize the batching function of the <code>autobot_vectorize</code> function?</h2>
<hr>
<p>I’ve tried these…</p>
<h3>
<a name="attempt-1-not-fully-utilizing-batching-in-autobot_vectorize-2" class="anchor" href="#attempt-1-not-fully-utilizing-batching-in-autobot_vectorize-2"></a>Attempt 1: Not fully utilizing batching in <code>autobot_vectorize</code>
</h3>
<p>I’ve tried doing this nesting the mapping function and by sending <code>autobot_vectorize</code> one input at a time but this is not fully utilizing its batching function capabilities, i.e.</p>
<pre><code class="lang-auto">
def optimus_prime(row):
    """This functions returns two data points with some arbitrary vectors.
    &gt;&gt;&gt; row = {'imagefile': 'train/0/16585.png', 'label': 0}
    &gt;&gt;&gt; optimus_prime(row)
    (tensor([1.23, 4.56, 7.89]), 0)
    """
    return autobot_vectorize(row['imagefile'])[0], row['label'] 
</code></pre>
<p>[Working snippet]:</p>
<pre><code class="lang-auto">import torch 

from torch.utils.data import DataLoader2
import pytorch_lightning as pl
from torchdata.datapipes.iter import IterDataPipe, IterableWrapper

content = """imagefile,label
train/0/16585.png,0
train/0/56789.png,0
train/0/47215.png,1
train/0/93155.png,1"""

with open('myfile.csv', 'w') as fout:
    fout.write(content)
    
def autobot_vectorize(imgfiles):
    # This vectorizer takes N imgfiles and return the 
    # a tensor of size N x 3.
    # For example, we are simulating this with torch.rand(N, 3)
    """
    &gt;&gt;&gt; autobot_vectorize(['train/0/47215.png', 'train/0/56789.png'])
    tensor([[1.23, 4.56, 7.89], [9.99, 8.88, 7.77]])
    """
    return torch.rand(len(imgfiles), 3)


def optimus_prime(row):
    """This functions returns two data points with some arbitrary vectors.
    &gt;&gt;&gt; row = {'imagefile': 'train/0/16585.png', 'label': 0}
    &gt;&gt;&gt; optimus_prime(row)
    (tensor([1.23, 4.56, 7.89]), 0)
    """
    return autobot_vectorize(row['imagefile'])[0], row['label'] 
    

class MyDataPipe(pl.LightningDataModule):
    def __init__(
        self,
        csv_files,
        skip_lines=0
    ):
        super().__init__()
        self.csv_files: list[str] = csv_files
        self.skip_lines: int = skip_lines

        # Initialize a datapipe.
        self.dp_chained_datapipe: IterDataPipe = (
            IterableWrapper(iterable=self.csv_files)
            .open_files()
            .parse_csv_as_dict(skip_lines=self.skip_lines)
        )
            
        self.dp_chained_datapipe = self.dp_chained_datapipe.map(optimus_prime)

    def train_dataloader(self, batch_size=1) -&gt; DataLoader2:
        return DataLoader2(dataset=self.dp_chained_datapipe, batch_size=batch_size)

dp = MyDataPipe(['myfile.csv'])

for row in dp.dp_chained_datapipe:
    print(row)
</code></pre>
<p>[out]:</p>
<pre><code class="lang-auto">(tensor([0.0740, 0.6826, 0.1215]), '0')
(tensor([0.0084, 0.2512, 0.8702]), '0')
(tensor([0.7141, 0.4003, 0.7872]), '1')
(tensor([0.8011, 0.4163, 0.9410]), '1')
</code></pre>
<h3>
<a name="attempt-2-batching-before-flatmap-but-cant-get-labels-to-be-included-3" class="anchor" href="#attempt-2-batching-before-flatmap-but-cant-get-labels-to-be-included-3"></a>Attempt 2: Batching before flatmap, but can’t get labels to be included…</h3>
<p>I’ve tried to process the X efficiently by kind of using <code>.flatmap()</code> but I don’t even know whether I’m really using the <code>autobot_vectorize</code> batching capabilities or is it still processing the batch, one at a time.</p>
<p>[Working snippet]:</p>
<pre><code class="lang-auto">import torch 

from torch.utils.data import DataLoader2
import pytorch_lightning as pl
from torchdata.datapipes.iter import IterDataPipe, IterableWrapper

content = """imagefile,label
train/0/16585.png,0
train/0/56789.png,0
train/0/47215.png,1
train/0/93155.png,1"""

with open('myfile.csv', 'w') as fout:
    fout.write(content)
    
def autobot_vectorize(imgfiles):
    # This vectorizer takes N imgfiles and return the 
    # a tensor of size N x 3.
    # For example, we are simulating this with torch.rand(N, 3)
    """
    &gt;&gt;&gt; autobot_vectorize(['train/0/47215.png', 'train/0/56789.png'])
    tensor([[1.23, 4.56, 7.89], [9.99, 8.88, 7.77]])
    """
    return torch.rand(len(imgfiles), 3)


class MyDataPipe(pl.LightningDataModule):
    def __init__(
        self,
        csv_files,
        skip_lines=0
    ):
        super().__init__()
        self.csv_files: list[str] = csv_files
        self.skip_lines: int = skip_lines

        # Initialize a datapipe.
        self.dp_chained_datapipe: IterDataPipe = (
            IterableWrapper(iterable=self.csv_files)
            .open_files()
            .parse_csv_as_dict(skip_lines=self.skip_lines)
        )
            
        self.dp_chained_datapipe = self.dp_chained_datapipe.map(lambda x: [x['imagefile']])
        self.dp_chained_datapipe = self.dp_chained_datapipe.flatmap(autobot_vectorize)

    def train_dataloader(self, batch_size=1) -&gt; DataLoader2:
        return DataLoader2(dataset=self.dp_chained_datapipe, batch_size=batch_size)

dp = MyDataPipe(['myfile.csv'])

for row in dp.dp_chained_datapipe:
    print(row)
</code></pre>
<p>Also, I’m not sure how to add the Ys (labels) when using flatmap as above.</p>
<h2>
<a name="what-is-the-a-simple-way-to-create-the-dataloader-from-the-datapipe-that-best-optimize-the-batching-function-of-the-autobot_vectorize-function-4" class="anchor" href="#what-is-the-a-simple-way-to-create-the-dataloader-from-the-datapipe-that-best-optimize-the-batching-function-of-the-autobot_vectorize-function-4"></a>What is the a simple way to create the DataLoader from the DataPipe that best optimize the batching function of the <code>autobot_vectorize</code> function?</h2>

alvations: <p>Also asked on <a href="https://stackoverflow.com/q/74046117/610569" class="inline-onebox" rel="noopener nofollow ugc">python - How to create DataPipe that best optimize the map-transform function that supports batching? - Stack Overflow</a></p>

nivek: <p>You can write your own IterDataPipe that collects <code>b</code> number of samples before transforming them and yielding that. It will look something like:</p>
<pre><code class="lang-python">def __iter__(self):
    img_batch = []
    label_batch = []
    for img, label in self.source_datapipe:
        img_batch.append(img)
        label_batch.append(label)
        if len(img_batch) == b:  # reaches batch size
            return autobot_vectorize(img_batch), process_labels(label_batch)
        # Extra logic to handle `drop_last` and etc
</code></pre>
<p>You will need a function to process the label since I am not sure what your desired output is for labels.</p>
<p>You can also do <code>dp.batch(batch_size).map(process_batch)</code> by defining a function that can handle <code>[(img1, label1), (img2, label2), (img3, label3), ...]</code> and create the desired output.</p>