00b_core.data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#default_exp core.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 00_core.data\n",
    "> Contains the original `DataLoader` from `fastai2`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from torch import multiprocessing\n",
    "from fastinference_pytorch.utils import tensor\n",
    "from numpy import ndarray\n",
    "from torch import Tensor\n",
    "from collections.abc import Sequence\n",
    "from torch.utils.data._utils.collate import default_collate, default_convert\n",
    "from torch.utils.data._utils.worker import get_worker_info\n",
    "from fastcore.all import *\n",
    "import typing, torch\n",
    "from torch.utils.data.dataloader import _MultiProcessingDataLoaderIter,_SingleProcessDataLoaderIter,_DatasetKind\n",
    "_loaders = (_MultiProcessingDataLoaderIter,_SingleProcessDataLoaderIter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import string\n",
    "bs = 4\n",
    "letters = list(string.ascii_lowercase)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def set_num_threads(nt):\n",
    "    \"Get numpy (and others) to use `nt` threads\"\n",
    "    try: import mkl; mkl.set_num_threads(nt)\n",
    "    except: pass\n",
    "    torch.set_num_threads(1)\n",
    "    os.environ['IPC_ENABLE']='1'\n",
    "    for o in ['OPENBLAS_NUM_THREADS','NUMEXPR_NUM_THREADS','OMP_NUM_THREADS','MKL_NUM_THREADS']:\n",
    "        os.environ[o] = str(nt)\n",
    "        \n",
    "def set_seed(s):\n",
    "    \"Set random seed for `random`, `torch`, and `numpy` (where available)\"\n",
    "    try: torch.manual_seed(s)\n",
    "    except NameError: pass\n",
    "    try: np.random.seed(s%(2**32-1))\n",
    "    except NameError: pass\n",
    "    random.seed(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def _wif(worker_id):\n",
    "    set_num_threads(1)\n",
    "    info = get_worker_info()\n",
    "    ds = info.dataset.d\n",
    "    ds.nw,ds.offs = info.num_workers,info.id\n",
    "    set_seed(info.seed)\n",
    "    ds.wif()\n",
    "\n",
    "class _FakeLoader:\n",
    "    _IterableDataset_len_called,_auto_collation,collate_fn,drop_last,dataset_kind,_dataset_kind,_index_sampler,generator,prefetch_factor = (\n",
    "        None,False,noops,False,_DatasetKind.Iterable,_DatasetKind.Iterable,Inf.count,None,2)\n",
    "    def __init__(self, d, pin_memory, num_workers, timeout):\n",
    "        self.dataset,self.default,self.worker_init_fn = self,d,_wif\n",
    "        store_attr(self, 'd,pin_memory,num_workers,timeout')\n",
    "\n",
    "    def __iter__(self): return iter(self.d.create_batches(self.d.sample()))\n",
    "\n",
    "    @property\n",
    "    def multiprocessing_context(self): return (None,multiprocessing)[self.num_workers>0]\n",
    "\n",
    "    @contextmanager\n",
    "    def no_multiproc(self):\n",
    "        old_nw = self.num_workers\n",
    "        try:\n",
    "            self.num_workers = 0\n",
    "            yield self.d\n",
    "        finally: self.num_workers = old_nw\n",
    "\n",
    "_collate_types = (ndarray, Tensor, typing.Mapping, str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def fa_collate(t):\n",
    "    b = t[0]\n",
    "    return (default_collate(t) if isinstance(b, _collate_types)\n",
    "            else type(t[0])([fa_collate(s) for s in zip(*t)]) if isinstance(b, Sequence)\n",
    "            else default_collate(t))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#e.g. x is int, y is tuple\n",
    "t = [(1,(2,3)),(1,(2,3))]\n",
    "test_eq(fa_collate(t), default_collate(t))\n",
    "test_eq(L(fa_collate(t)).map(type), [Tensor,tuple])\n",
    "\n",
    "t = [(1,(2,(3,4))),(1,(2,(3,4)))]\n",
    "test_eq(fa_collate(t), default_collate(t))\n",
    "test_eq(L(fa_collate(t)).map(type), [Tensor,tuple])\n",
    "test_eq(L(fa_collate(t)[1]).map(type), [Tensor,tuple])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def fa_convert(t):\n",
    "    return (default_convert(t) if isinstance(t, _collate_types)\n",
    "            else type(t)([fa_convert(s) for s in t]) if isinstance(t, Sequence)\n",
    "            else default_convert(t))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t0 = array([1,2])\n",
    "t = [t0,(t0,t0)]\n",
    "\n",
    "test_eq(fa_convert(t), default_convert(t))\n",
    "test_eq(L(fa_convert(t)).map(type), [Tensor,tuple])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SkipItemException(Exception): pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@funcs_kwargs\n",
    "class DataLoader(GetAttr):\n",
    "    _noop_methods = 'wif before_iter after_item before_batch after_batch after_iter'.split()\n",
    "    for o in _noop_methods:\n",
    "        exec(f\"def {o}(self, x=None, *args, **kwargs): return x\")\n",
    "    _methods = _noop_methods + 'create_batches create_item create_batch retain \\\n",
    "        get_idxs sample shuffle_fn do_batch create_batch'.split()\n",
    "    _default = 'dataset'\n",
    "    def __init__(self, dataset=None, bs=None, num_workers=0, pin_memory=False, timeout=0, batch_size=None,\n",
    "                 shuffle=False, drop_last=False, indexed=None, n=None, device=None, **kwargs):\n",
    "        if batch_size is not None: bs = batch_size # PyTorch compatibility\n",
    "        assert not (bs is None and drop_last)\n",
    "        if indexed is None: indexed = dataset is not None and hasattr(dataset,'__getitem__')\n",
    "        if n is None:\n",
    "            try: n = len(dataset)\n",
    "            except TypeError: pass\n",
    "        store_attr(self, 'dataset,bs,shuffle,drop_last,indexed,n,pin_memory,timeout,device')\n",
    "        self.rng,self.nw,self.offs = random.Random(random.randint(0,2**32-1)),1,0\n",
    "        self.fake_l = _FakeLoader(self, pin_memory, num_workers, timeout)\n",
    "\n",
    "    def __len__(self):\n",
    "        if self.n is None: raise TypeError\n",
    "        if self.bs is None: return self.n\n",
    "        return self.n//self.bs + (0 if self.drop_last or self.n%self.bs==0 else 1)\n",
    "\n",
    "    def get_idxs(self):\n",
    "        idxs = Inf.count if self.indexed else Inf.nones\n",
    "        if self.n is not None: idxs = list(itertools.islice(idxs, self.n))\n",
    "        if self.shuffle: idxs = self.shuffle_fn(idxs)\n",
    "        return idxs\n",
    "\n",
    "    def sample(self):\n",
    "        idxs = self.get_idxs()\n",
    "        return (b for i,b in enumerate(idxs) if i//(self.bs or 1)%self.nw==self.offs)\n",
    "\n",
    "    def __iter__(self):\n",
    "        self.randomize()\n",
    "        self.before_iter()\n",
    "        for b in _loaders[self.fake_l.num_workers==0](self.fake_l):\n",
    "            if self.device is not None: b = to_device(b, self.device)\n",
    "            yield self.after_batch(b)\n",
    "        self.after_iter()\n",
    "        if hasattr(self, 'it'): delattr(self, 'it')\n",
    "\n",
    "    def create_batches(self, samps):\n",
    "        self.it = iter(self.dataset) if self.dataset is not None else None\n",
    "        res = filter(lambda o:o is not None, map(self.do_item, samps))\n",
    "        yield from map(self.do_batch, self.chunkify(res))\n",
    "\n",
    "    def new(self, dataset=None, cls=None, **kwargs):\n",
    "        if dataset is None: dataset = self.dataset\n",
    "        if cls is None: cls = type(self)\n",
    "        cur_kwargs = dict(dataset=dataset, num_workers=self.fake_l.num_workers, pin_memory=self.pin_memory, timeout=self.timeout,\n",
    "                          bs=self.bs, shuffle=self.shuffle, drop_last=self.drop_last, indexed=self.indexed, device=self.device)\n",
    "        for n in self._methods: cur_kwargs[n] = getattr(self, n)\n",
    "        return cls(**merge(cur_kwargs, kwargs))\n",
    "\n",
    "    @property\n",
    "    def prebatched(self): return self.bs is None\n",
    "    def do_item(self, s):\n",
    "        try: return self.after_item(self.create_item(s))\n",
    "        except SkipItemException: return None\n",
    "    def chunkify(self, b): return b if self.prebatched else chunked(b, self.bs, self.drop_last)\n",
    "    def shuffle_fn(self, idxs): return self.rng.sample(idxs, len(idxs))\n",
    "    def randomize(self): self.rng = random.Random(self.rng.randint(0,2**32-1))\n",
    "    def retain(self, res, b):  return retain_types(res, b[0] if is_listy(b) else b)\n",
    "    def create_item(self, s):  return next(self.it) if s is None else self.dataset[s]\n",
    "    def create_batch(self, b): return (fa_collate,fa_convert)[self.prebatched](b)\n",
    "    def do_batch(self, b): return self.retain(self.create_batch(self.before_batch(b)), b)\n",
    "    def to(self, device): self.device = device\n",
    "    def one_batch(self):\n",
    "        if self.n is not None and len(self)==0: raise ValueError(f'This DataLoader does not contain any batches')\n",
    "        with self.fake_l.no_multiproc(): res = first(self)\n",
    "        if hasattr(self, 'it'): delattr(self, 'it')\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(#22) [0.5807299833206229,0.7357323148365269,0.28225862923257816,0.9424049519160917,0.15428878174667815,0.9115489115549179,0.5549914373946663,0.19098497518115942,0.7453394325264376,0.4431252777626504...]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class RandDL(DataLoader):\n",
    "    def create_item(self, s):\n",
    "        r = random.random()\n",
    "        return r if r<0.95 else stop()\n",
    "\n",
    "L(RandDL())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(#3) [4,4,4]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "L(RandDL(bs=4, drop_last=True)).map(len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(#9) [4,4,4,4,4,4,4,4,4]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dl = RandDL(bs=4, num_workers=4, drop_last=True)\n",
    "L(dl).map(len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(dl.fake_l.num_workers, 4)\n",
    "with dl.fake_l.no_multiproc(): \n",
    "    test_eq(dl.fake_l.num_workers, 0)\n",
    "    L(dl).map(len)\n",
    "test_eq(dl.fake_l.num_workers, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(#0) []"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _rand_item(s):\n",
    "    r = random.random()\n",
    "    return r if r<0.95 else stop()\n",
    "\n",
    "L(DataLoader(create_item=_rand_item))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds1 = DataLoader(letters)\n",
    "test_eq(L(ds1), letters)\n",
    "test_eq(len(ds1), 26)\n",
    "\n",
    "test_shuffled(L(DataLoader(letters, shuffle=True)), letters)\n",
    "\n",
    "ds1 = DataLoader(letters, indexed=False)\n",
    "test_eq(L(ds1), letters)\n",
    "test_eq(len(ds1), 26)\n",
    "\n",
    "t2 = L(tensor([0,1,2]),tensor([3,4,5]))\n",
    "ds2 = DataLoader(t2)\n",
    "test_eq_type(L(ds2), t2)\n",
    "\n",
    "t3 = L(array([0,1,2]),array([3,4,5]))\n",
    "ds3 = DataLoader(t3)\n",
    "test_eq_type(L(ds3), t3.map(tensor))\n",
    "\n",
    "ds4 = DataLoader(t3, create_batch=noop, after_iter=lambda: setattr(t3, 'f', 1))\n",
    "test_eq_type(L(ds4), t3)\n",
    "test_eq(t3.f, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TfmdDL(DataLoader):\n",
    "    \"Transformed `DataLoader`\"\n",
    "    def __init__(self, dataset, bs=64, shuffle=False, num_workers=None, verbose=False, do_setup=True, **kwargs):\n",
    "        if num_workers is None: num_workers = min(16, defaults.cpus)\n",
    "        for nm in _batch_tfms: kwargs[nm] = Pipeline(kwargs.get(nm,None))\n",
    "        super().__init__(dataset, bs=bs, shuffle=shuffle, num_workers=num_workers, **kwargs)\n",
    "        if do_setup:\n",
    "            for nm in _batch_tfms:\n",
    "                pv(f\"Setting up {nm}: {kwargs[nm]}\", verbose)\n",
    "                kwargs[nm].setup(self)\n",
    "\n",
    "    def _one_pass(self):\n",
    "        b = self.do_batch([self.do_item(0)])\n",
    "        if self.device is not None: b = to_device(b, self.device)\n",
    "        its = self.after_batch(b)\n",
    "        self._n_inp = 1 if not isinstance(its, (list,tuple)) or len(its)==1 else len(its)-1\n",
    "        self._types = explode_types(its)\n",
    "\n",
    "    def _retain_dl(self,b):\n",
    "        if not getattr(self, '_types', None): self._one_pass()\n",
    "        return retain_types(b, typs=self._types)\n",
    "\n",
    "    @delegates(DataLoader.new)\n",
    "    def new(self, dataset=None, cls=None, **kwargs):\n",
    "        res = super().new(dataset, cls, do_setup=False, **kwargs)\n",
    "        if not hasattr(self, '_n_inp') or not hasattr(self, '_types'):\n",
    "            try:\n",
    "                self._one_pass()\n",
    "                res._n_inp,res._types = self._n_inp,self._types\n",
    "            except: print(\"Could not do one pass in your dataloader, there is something wrong in it\")\n",
    "        else: res._n_inp,res._types = self._n_inp,self._types\n",
    "        return res\n",
    "\n",
    "    def before_iter(self):\n",
    "        super().before_iter()\n",
    "        split_idx = getattr(self.dataset, 'split_idx', None)\n",
    "        for nm in _batch_tfms:\n",
    "            f = getattr(self,nm)\n",
    "            if isinstance(f,Pipeline): f.split_idx=split_idx\n",
    "\n",
    "    def decode(self, b): return self.before_batch.decode(to_cpu(self.after_batch.decode(self._retain_dl(b))))\n",
    "    def decode_batch(self, b, max_n=9, full=True): return self._decode_batch(self.decode(b), max_n, full)\n",
    "\n",
    "    def _decode_batch(self, b, max_n=9, full=True):\n",
    "        f = self.after_item.decode\n",
    "        f = compose(f, partial(getattr(self.dataset,'decode',noop), full = full))\n",
    "        return L(batch_to_samples(b, max_n=max_n)).map(f)\n",
    "\n",
    "    @property\n",
    "    def n_inp(self):\n",
    "        if hasattr(self.dataset, 'n_inp'): return self.dataset.n_inp\n",
    "        if not hasattr(self, '_n_inp'): self._one_pass()\n",
    "        return self._n_inp\n",
    "\n",
    "    def to(self, device):\n",
    "        self.device = device\n",
    "        for tfm in self.after_batch.fs:\n",
    "            for a in L(getattr(tfm, 'parameters', None)): setattr(tfm, a, getattr(tfm, a).to(device))\n",
    "        return self"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataLoaders(GetAttr):\n",
    "    \"Basic wrapper around several `DataLoader`s.\"\n",
    "    _default='train'\n",
    "    def __init__(self, *loaders, path='.', device=None):\n",
    "        self.loaders,self.path = list(loaders),Path(path)\n",
    "        if device is not None or hasattr(loaders[0],'to'): self.device = device\n",
    "\n",
    "    def __getitem__(self, i): return self.loaders[i]\n",
    "    def new_empty(self):\n",
    "        loaders = [dl.new(dl.dataset.new_empty()) for dl in self.loaders]\n",
    "        return type(self)(*loaders, path=self.path, device=self.device)\n",
    "\n",
    "    def _set(i, self, v): self.loaders[i] = v\n",
    "    train   ,valid    = add_props(lambda i,x: x[i], _set)\n",
    "    train_ds,valid_ds = add_props(lambda i,x: x[i].dataset)\n",
    "\n",
    "    @property\n",
    "    def device(self): return self._device\n",
    "\n",
    "    @device.setter\n",
    "    def device(self, d):\n",
    "        for dl in self.loaders: dl.to(d)\n",
    "        self._device = d\n",
    "\n",
    "    def to(self, device):\n",
    "        self.device = device\n",
    "        return self\n",
    "\n",
    "    def cuda(self): return self.to(device=default_device())\n",
    "    def cpu(self):  return self.to(device=torch.device('cpu'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}