DataLoader does not access updated Dataset attribute when num_workers > 0

username: jgscott
cooked: <p>The context here is a scientific machine learning problem in which training data (which happens to be sequences) is simulated from an underlying ground-truth scientific model.  A key wrinkle is that it’s important that the model sees simulations of different sizes during training.  So a minibatch of data has shape [batch_size, seq_length, seq_depth], where seq_depth is always the same for a given problem, but seq_length needs to change from one minibatch to the next.</p>
<p>Because the simulation of the data is itself nontrivial, I want to use a Dataloader with multiple workers during training.  The issue I’m running into is that the workers don’t “see” a change I make to a Dataset attribute that varies the sequence length during training.</p>
<p>Here’s a minimal reproducible example where seq_depth=2 to illustrate the issue.</p>
<pre><code class="lang-auto">import torch
import numpy as np 
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader, Dataset

def my_gen_sample(n):
    return torch.randn(n,2)

class TestDataset(Dataset):
    def __init__(self, n_sims, gen_sample):
        self.n_sims = n_sims
        self.gen_sample = gen_sample
        self.seq_length = 1  # attribute to be modified after each minibatch

    def __len__(self):
        return self.n_sims

    def set_seq_length(self, n: int):
        self.seq_length = n

    def __getitem__(self, idx):
        x = self.gen_sample(self.seq_length)
        return x
</code></pre>
<p>This works as expected when num_workers=0.  But when num_workers &gt; 0, the changes to ds.seq_length after each minibatch are no longer reflected.</p>
<p>Code for num_workers=0:</p>
<pre><code class="lang-auto">ds = TestDataset(50*64, my_gen_sample)
dl = DataLoader(ds, batch_size=64, num_workers=0)

ds.set_seq_length(33)
lengths = [None]*ds.__len__()
for i, x_b in enumerate(dl):
    seq_length = np.random.randint(low=5, high=50, size=1)[0]
    ds.set_seq_length(seq_length)
    lengths[i] = x_b.shape[1]
</code></pre>
<p>Now compared with num_workers=2:</p>
<pre><code class="lang-auto">dl = DataLoader(ds, batch_size=64, num_workers=2)

ds.set_seq_length(33)
lengths2 = [None]*ds.__len__()
for i, x_b in enumerate(dl):
    seq_length = np.random.randint(low=5, high=50, size=1)[0]
    ds.set_seq_length(seq_length)
    lengths2[i] = x_b.shape[1]

plt.plot(lengths)
plt.plot(lengths2)
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/7/e74222c40ffbf31c41570a9064724a96012977b1.png" alt="image" data-base62-sha1="wZNZrkrzCAJFQWXZMrpLyJ2mGlz" width="543" height="413"></p>
<p>I would be grateful to any suggestions for how to get the num_workers = 0 behavior with multiple workers. Thank you!</p>

username: ptrblck
cooked: <p>The issue is expected since each worker is a separate process working on a copy of the <code>Dataset</code>.<br>
You could try to share a data structure between all workers and manipulate it from the main thread e.g. via:</p>
<pre data-code-wrap="python"><code class="lang-python">class TestDataset(Dataset):
    def __init__(self, n_sims, gen_sample, num_workers=0):
        self.n_sims = n_sims
        self.gen_sample = gen_sample
        shared_array_base = mp.Array(ctypes.c_int64, 1)
        shared_array = np.ctypeslib.as_array(shared_array_base.get_obj())
        self.seq_length = torch.from_numpy(shared_array)
</code></pre>
<p>but you would run into race conditions since each worker loads and processes its batch independently.</p>
<aside class="quote no-group" data-username="jgscott" data-post="1" data-topic="176879">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/958977/48.png" class="avatar"> jgscott:</div>
<blockquote>
<p>So a minibatch of data has shape [batch_size, seq_length, seq_depth], where seq_depth is always the same for a given problem, but seq_length needs to change from one minibatch to the next.</p>
</blockquote>
</aside>
<p>Based on this description it sounds as if you could know the sequence length in advance for each batch. If so you might be able to write a custom <code>BatchSampler</code> instead which allows you to pass the indices and sequence length of the entire batch to the <code>__getitem__</code> as seen in <a href="https://discuss.pytorch.org/t/how-to-use-batchsampler-with-getitem-dataset/78788/4">this post</a>.</p>
<p>Something like this might work for you:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import numpy as np 
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader, Dataset


def my_gen_sample(n):
    return torch.randn(n,2)


class TestDataset(Dataset):
    def __init__(self, n_sims, gen_sample, num_workers=0):
        self.n_sims = n_sims
        self.gen_sample = gen_sample
        self.seq_length = 1  # attribute to be modified after each minibatch

    def __len__(self):
        return self.n_sims

    def set_seq_length(self, n: int):
        self.seq_length = n

    def __getitem__(self, idx):
        x = self.gen_sample(self.seq_length)
        return x
    
    
ds = TestDataset(50*64, my_gen_sample)
dl = DataLoader(ds, batch_size=64, num_workers=0)

ds.set_seq_length(33)
lengths = [None] * len(dl)
for i, x_b in enumerate(dl):
    seq_length = np.random.randint(low=5, high=50, size=1)[0]
    ds.set_seq_length(seq_length)
    lengths[i] = x_b.shape[1]


class MyBatchSampler(torch.utils.data.sampler.Sampler):
    def __init__(self, data_indices, seq_lengths, batch_size):
        self.data_indices = data_indices
        self.seq_lengths = seq_lengths
        self.batch_size = batch_size

    def __iter__(self):
        data_indices_iter = iter(self.data_indices)
        seq_lengths_iter = iter(self.seq_lengths)
        while True:
            try:
                batch = [(next(data_indices_iter), next(seq_lengths_iter)) for _ in range(self.batch_size)]
                yield batch
            except StopIteration:
                break
        
    def __len__(self) -&gt; int:
        return len(self.data_indices) // self.batch_size  # type: ignore[arg-type]


class MyTestDataset(Dataset):
    def __init__(self, n_sims, gen_sample, num_workers=0):
        self.n_sims = n_sims
        self.gen_sample = gen_sample

    def __len__(self):
        return self.n_sims

    def __getitem__(self, idx):
        data_indices, seq_lengths = list(zip(*idx))
        x = []
        for data_index, seq_len in zip(data_indices, seq_lengths):
            x.append(self.gen_sample(seq_len))
        x = torch.stack(x)
        return x


ds = MyTestDataset(50*64, my_gen_sample)    
batch_size = 64
data_indices = torch.arange(len(ds))
seq_lengths = torch.tensor(lengths[:50]).repeat_interleave(batch_size)
assert len(data_indices) == len(seq_lengths)
sampler = MyBatchSampler(data_indices, seq_lengths, batch_size)

dl = DataLoader(ds, sampler=sampler, num_workers=2)

lengths2 = [None] * len(dl)
for i, x_b in enumerate(dl):
    x_b = x_b.squeeze(0)
    lengths2[i] = x_b.shape[1]

plt.plot(lengths)
plt.plot(lengths2)

print((torch.tensor(lengths) == torch.tensor(lengths2)).all())
# tensor(True)
</code></pre>

username: jgscott
cooked: <p>Thank you so much!  This seems like a perfect solution.  I had wondered about a solution where the Dataset itself stored the lengths in your MyBatchSampler.seq_lengths(), so that a copy of which idx’s went with which sequence lengths was copied to all workers in a manner that ensure consistent seq_length within a single batch.  But that required the Dataset to know about what batch size to expect during training.  Your solution seems to encapsulate those sequence lengths more straightforwardly.</p>

username: ptrblck
cooked: <p>Great to hear the approach might work for you.<br>
Just wanted to let you know that the <code>num_workers</code> argument in <code>TestDataset</code> is dead code from a previous test, so ignore it. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>