/usr/local/lib/python3.8/dist-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0

bs = 4
letters = list(string.ascii_lowercase)

DataLoader helpers

fastai includes a replacement for Pytorch's DataLoader which is largely API-compatible, and adds a lot of useful functionality and flexibility. Before we look at the class, there are a couple of helpers we'll need to define.

fa_collate[source]

fa_collate(t)

A replacement for PyTorch default_collate which maintains types and handles Sequences

t = [(1,(2,3)),(1,(2,3))]
test_eq(fa_collate(t), default_collate(t))
test_eq(L(fa_collate(t)).map(type), [Tensor,tuple])

t = [(1,(2,(3,4))),(1,(2,(3,4)))]
test_eq(fa_collate(t), default_collate(t))
test_eq(L(fa_collate(t)).map(type), [Tensor,tuple])
test_eq(L(fa_collate(t)[1]).map(type), [Tensor,tuple])

fa_convert[source]

fa_convert(t)

A replacement for PyTorch default_convert which maintains types and handles Sequences

t0 = array([1,2])
t = [t0,(t0,t0)]

test_eq(fa_convert(t), default_convert(t))
test_eq(L(fa_convert(t)).map(type), [Tensor,tuple])

class SkipItemException[source]

SkipItemException() :: Exception

Raised to notify DataLoader to skip an item

class DataLoader[source]

DataLoader(dataset=None, bs=None, num_workers=0, pin_memory=False, timeout=0, batch_size=None, shuffle=False, drop_last=False, indexed=None, n=None, device=None, persistent_workers=False, wif=None, before_iter=None, after_item=None, before_batch=None, after_batch=None, after_iter=None, create_batches=None, create_item=None, create_batch=None, retain=None, get_idxs=None, sample=None, shuffle_fn=None, do_batch=None) :: GetAttr

API compatible with PyTorch DataLoader, with a lot more callbacks and flexibility

Arguments to DataLoader:

• dataset: dataset from which to load the data. Can be either map-style or iterable-style dataset.
• bs (int): how many samples per batch to load (if batch_size is provided then batch_size will override bs). If bs=None, then it is assumed that dataset.__getitem__ returns a batch.
• num_workers (int): how many subprocesses to use for data loading. 0 means that the data will be loaded in the main process.
• pin_memory (bool): If True, the data loader will copy Tensors into CUDA pinned memory before returning them.
• timeout (float>0): the timeout value in seconds for collecting a batch from workers.
• batch_size (int): It is only provided for PyTorch compatibility. Use bs.
• shuffle (bool): If True, then data is shuffled every time dataloader is fully read/iterated.
• drop_last (bool): If True, then the last incomplete batch is dropped.
• indexed (bool): The DataLoader will make a guess as to whether the dataset can be indexed (or is iterable), but you can override it with this parameter. True by default.
• n (int): Defaults to len(dataset). If you are using iterable-style dataset, you can specify the size with n.
• device (torch.device): Defaults to default_device() which is CUDA by default. You can specify device as torch.device('cpu').

Override item and use the default infinite sampler to get a stream of unknown length (stop() when you want to stop the stream).

class RandDL(DataLoader):
    def create_item(self, s):
        r = random.random()
        return r if r<0.95 else stop()

L(RandDL())

(#80) [0.7080409288034344,0.03380592302379748,0.3295210988884517,0.3984895442219716,0.2466406073732288,0.784596719349558,0.7405184556807134,0.8781785423004165,0.06125487321640455,0.49829694909644495...]

L(RandDL(bs=4, drop_last=True)).map(len)

(#2) [4,4]

dl = RandDL(bs=4, num_workers=4, drop_last=True)
L(dl).map(len)

(#20) [4,4,4,4,4,4,4,4,4,4...]

test_num_workers = 0 if sys.platform == "win32" else 4
test_eq(dl.fake_l.num_workers, test_num_workers)
with dl.fake_l.no_multiproc(): 
    test_eq(dl.fake_l.num_workers, 0)
    L(dl).map(len)
test_eq(dl.fake_l.num_workers, test_num_workers)

def _rand_item(s):
    r = random.random()
    return r if r<0.95 else stop()

L(DataLoader(create_item=_rand_item))

(#50) [0.013926194175715834,0.5455661909308923,0.5540701885046865,0.38608142489299,0.11039197023987835,0.5609994837228025,0.511809356329029,0.21402937998255644,0.2856955111775398,0.3976737532547229...]

If you don't set bs, then dataset is assumed to provide an iterator or a __getitem__ that returns a batch.

ds1 = DataLoader(letters)
test_eq(L(ds1), letters)
test_eq(len(ds1), 26)

test_shuffled(L(DataLoader(letters, shuffle=True)), letters)

ds1 = DataLoader(letters, indexed=False)
test_eq(L(ds1), letters)
test_eq(len(ds1), 26)

t2 = L(tensor([0,1,2]),tensor([3,4,5]))
ds2 = DataLoader(t2)
test_eq_type(L(ds2), t2)

t3 = L(array([0,1,2], dtype=np.int64),array([3,4,5], dtype=np.int64))
ds3 = DataLoader(t3)
test_eq_type(L(ds3), t3.map(tensor))

ds4 = DataLoader(t3, create_batch=noop, after_iter=lambda: setattr(t3, 'f', 1))
test_eq_type(L(ds4), t3)
test_eq(t3.f, 1)

If you do set bs, then dataset is assumed to provide an iterator or a __getitem__ that returns a single item of a batch.

def twoepochs(d): return ' '.join(''.join(list(o)) for _ in range(2) for o in d)

ds1 = DataLoader(letters, bs=4, drop_last=True, num_workers=0)
test_eq(twoepochs(ds1), 'abcd efgh ijkl mnop qrst uvwx abcd efgh ijkl mnop qrst uvwx')

ds1 = DataLoader(letters,4,num_workers=2)
test_eq(twoepochs(ds1), 'abcd efgh ijkl mnop qrst uvwx yz abcd efgh ijkl mnop qrst uvwx yz')

ds1 = DataLoader(range(12), bs=4, num_workers=3)
test_eq_type(L(ds1), L(tensor([0,1,2,3]),tensor([4,5,6,7]),tensor([8,9,10,11])))

ds1 = DataLoader([str(i) for i in range(11)], bs=4, after_iter=lambda: setattr(t3, 'f', 2))
test_eq_type(L(ds1), L(['0','1','2','3'],['4','5','6','7'],['8','9','10']))
test_eq(t3.f, 2)

it = iter(DataLoader(map(noop,range(20)), bs=4, num_workers=1))
test_eq_type([next(it) for _ in range(3)], [tensor([0,1,2,3]),tensor([4,5,6,7]),tensor([8,9,10,11])])

Iterable dataloaders require specific tests.

class DummyIterableDataset(IterableDataset):
    def __iter__(self):
        yield from range(11)

ds1 = DataLoader(DummyIterableDataset(), bs=4)
# Check it yields fine, and check we can do multiple passes
for i in range(3):
    test_eq_type(L(ds1), L(tensor([0,1,2,3]),tensor([4,5,6,7]),tensor([8,9,10])))

# Check `drop_last` works fine (with multiple passes, since this will prematurely terminate the iterator)
ds1 = DataLoader(DummyIterableDataset(), bs=4, drop_last=True)
for i in range(3):
    test_eq_type(L(ds1), L(tensor([0,1,2,3]),tensor([4,5,6,7])))

class SleepyDL(list):
    def __getitem__(self,i):
        time.sleep(random.random()/50)
        return super().__getitem__(i)

t = SleepyDL(letters)

%time test_eq(DataLoader(t, num_workers=0), letters)
%time test_eq(DataLoader(t, num_workers=2), letters)
%time test_eq(DataLoader(t, num_workers=4), letters)

dl = DataLoader(t, shuffle=True, num_workers=1)
test_shuffled(L(dl), letters)
test_shuffled(L(dl), L(dl))
L(dl)

CPU times: user 4 ms, sys: 0 ns, total: 4 ms
Wall time: 268 ms
CPU times: user 12 ms, sys: 28 ms, total: 40 ms
Wall time: 196 ms
CPU times: user 32 ms, sys: 20 ms, total: 52 ms
Wall time: 137 ms

(#26) ['b','z','j','i','c','m','l','p','u','v'...]

class SleepyQueue():
    "Simulate a queue with varying latency"
    def __init__(self, q): self.q=q
    def __iter__(self):
        while True:
            time.sleep(random.random()/100)
            try: yield self.q.get_nowait()
            except queues.Empty: return

q = Queue()
for o in range(30): q.put(o)
it = SleepyQueue(q)

if not (sys.platform == "win32" and IN_NOTEBOOK):
    %time test_shuffled(L(DataLoader(it, num_workers=4)), L(range(30)))

CPU times: user 8 ms, sys: 44 ms, total: 52 ms
Wall time: 110 ms

class A(TensorBase): pass

for nw in (0,2):
    t = A(tensor([1,2]))
    dl = DataLoader([t,t,t,t,t,t,t,t], bs=4, num_workers=nw)
    b = first(dl)
    test_eq(type(b), A)

    t = (A(tensor([1,2])),)
    dl = DataLoader([t,t,t,t,t,t,t,t], bs=4, num_workers=nw)
    b = first(dl)
    test_eq(type(b[0]), A)

list(DataLoader(list(range(50)),bs=32,shuffle=True,num_workers=3))

[tensor([24, 25, 15, 38, 36, 42,  6, 18, 17,  3, 47, 33, 22, 44, 11, 23,  9, 49,
         31, 30,  8, 37, 14, 41, 27, 13, 34,  1, 12, 21,  2, 28]),
 tensor([45, 39, 43, 26,  7, 35, 32, 20, 46, 48, 40, 19,  5, 16, 10,  0, 29,  4])]

class A(TensorBase): pass
t = A(tensor(1,2))

tdl = DataLoader([t,t,t,t,t,t,t,t], bs=4, num_workers=2, after_batch=to_device)
b = first(tdl)
test_eq(type(b), A)

# Unknown attributes are delegated to `dataset`
test_eq(tdl.pop(), tensor(1,2))

Override get_idxs to return the same index until consumption of the DL. This is intented to test consistent sampling behavior when num_workers>1.

class AdamantDL(DataLoader):
    def get_idxs(self):
        r=random.randint(0,self.n-1)
        return [r] * self.n

test_eq(torch.cat(tuple(AdamantDL((list(range(50))),bs=16,num_workers=4))).unique().numel(),1)

# # test num_workers > 0 in scripts works when python process start method is spawn
# process = Popen(["python", "dltest.py"], stdout=PIPE)
# _, err = process.communicate(timeout=15)
# exit_code = process.wait()
# test_eq(exit_code, 0)