mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-04-28 11:45:23 +00:00
19d1510ea2
* support auto distributed data loader * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * support auto distributed data loader * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix tp error * remove unused parameters * remove unused * update inference * update docs * update inference --------- Co-authored-by: Michelle <qianranma8@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
39 lines
949 B
Python
39 lines
949 B
Python
from abc import abstractstaticmethod
|
|
|
|
from colossal_eval.utils import jdump
|
|
from torch.utils.data import Dataset
|
|
|
|
from colossalai.logging import DistributedLogger
|
|
|
|
|
|
class BaseDataset:
|
|
"""
|
|
Base class for dataset wrapper.
|
|
|
|
Args:
|
|
path: The path to the original dataset.
|
|
logger: Logger for the dataset.
|
|
"""
|
|
|
|
def __init__(self, path, logger, *args, **kwargs):
|
|
self.dataset = self.load(path, logger, *args, **kwargs)
|
|
|
|
def save(self, save_path):
|
|
"""Save the converted dataset"""
|
|
jdump(self.dataset, save_path)
|
|
|
|
@abstractstaticmethod
|
|
def load(path, logger: DistributedLogger, *args, **kwargs):
|
|
"""Load the original dataset and convert it into the inference dataset"""
|
|
|
|
|
|
class DistributedDataset(Dataset):
|
|
def __init__(self, data):
|
|
self.data = data
|
|
|
|
def __len__(self):
|
|
return len(self.data)
|
|
|
|
def __getitem__(self, idx):
|
|
return self.data[idx]
|