下面给出的代码，允许用户使用LMDataset对象，加快数据集的访问速度。它预先读取传入dataset中的数据，并存储于LMDB数据库中。在ImageNet的测试表明，它能够加快图像读取速度4.25倍。

使用代码如下：

from LMDataset import LMDataset
from torchvision.datasets import ImageNet
 
 
if __name__ == "__main__":
    print("Begin Caching")
    dataset_train: Dataset[torch.Tensor] = ImageNet(
        "~/dataset/ImageNet-1k", split="train"
    )
    cached_dataset = LMDataset(dataset_train, "~/dataset/ImageNet-1k", "train")
 
    for X, y in cached_dataset:
        # ...

源码LMDataset.py如下：

# -*- coding: utf-8 -*-
import lmdb
import pickle
from tqdm import tqdm
import multiprocessing as mp
from os import path, makedirs
from torch.utils.data import Dataset, DataLoader
from typing import TypeVar, Optional, Literal, Sized, Tuple, Callable
from torch.nn import Identity
 
__all__ = ["LMDataset"]
 
# Begin Configurations
MAX_SIZE = 1024**4  # 1TB
NUM_WORKERS = max(8, mp.cpu_count())
MAX_READERS = max(128, 2 * mp.cpu_count())
PICKLE_PROTOCOL = pickle.HIGHEST_PROTOCOL
INDEX_SIZE = 4
BATCH_SIZE = 256
# End Configurations
 
 
dtype = TypeVar("dtype", covariant=True)
transform_type = TypeVar("transform_type", covariant=True)
 
 
class _PickleWrapper(Dataset[bytes]):
    def __init__(self, dataset: Dataset[dtype]) -> None:
        super().__init__()
        assert isinstance(dataset, Sized)
        self.dataset = dataset
 
    def __len__(self) -> int:
        return len(self.dataset)
 
    def __getitem__(self, index: int) -> Tuple[bytes, bytes]:
        value = pickle.dumps(self.dataset[index], protocol=PICKLE_PROTOCOL)
        key = index.to_bytes(INDEX_SIZE, "little")
        return key, value
 
 
class LMDataset(Dataset[transform_type]):
    def __init__( 
        self,
        dataset: Optional[Dataset[dtype]],
        root: str,
        split: Literal["train", "val", "test"],
        transform: Callable[[dtype], transform_type] = Identity(),
        desc: Optional[str] = "Caching Dataset",
    ) -> None:
        assert isinstance(dataset, Sized)
        self.root = path.join(path.expanduser(root), ".LMDB", f"{split}.mdb")
        self.dataset = dataset
        makedirs(path.dirname(self.root), exist_ok=True)
 
        # Cache Dataset if not already cached
        if not path.isfile(self.root):
            if dataset is not None:
                self._cache_dataset(dataset, desc=desc)
 
        # Open Read-only LMDB Environment
        self.env = lmdb.Environment(
            self.root,
            map_size=MAX_SIZE,
            subdir=False,
            readonly=True,
            max_readers=MAX_READERS,
            max_dbs=4,
            lock=False,
        )
 
        # Check if dataset is complete
        self.length = self._len_db()
        if dataset is not None:
            assert len(dataset) == self.length, "Dataset Length Mismatch!"
 
        self.transform = transform
 
    def _cache_dataset(
        self, dataset: Dataset[dtype], desc: Optional[str] = None
    ) -> None:
        assert isinstance(dataset, Sized)
        makedirs(path.dirname(self.root), exist_ok=True)
 
        pickle_dataset = _PickleWrapper(dataset)
 
        data_loader = DataLoader(
            pickle_dataset,
            batch_size=BATCH_SIZE,
            num_workers=NUM_WORKERS,
            prefetch_factor=8,
            persistent_workers=True,
        )
 
        if desc is not None:
            data_loader = tqdm(data_loader, desc=desc)
 
        with lmdb.Environment(
            self.root,
            map_size=MAX_SIZE,
            subdir=False,
            readonly=False,
            max_readers=1,
            readahead=False,
            sync=False,
            map_async=True,
            metasync=False,
            meminit=False,
            lock=True,
        ) as write_env:
            with write_env.begin(write=True, buffers=True) as txn:
                for batch in data_loader:
                    for key, value in zip(*batch):
                        txn.put(key, value)
 
    def __len__(self) -> int:
        return self.length
 
    def __getitem__(self, index: int) -> transform_type:
        key = index.to_bytes(INDEX_SIZE, "little")
        with self.env.begin(write=False, buffers=True) as txn:
            value: bytes = txn.get(key)  # type: ignore
        return self.transform(pickle.loads(value))
 
    def _len_db(self) -> int:
        with self.env.begin(write=False, buffers=True) as txn:
            stat = txn.stat(db=None)
        return stat["entries"]
 
    def __getattr__(self, name: str):
        return getattr(self.dataset, name, getattr(self.dataset, name))