HuggingfaceのDataLoaderとDatacollatorのソースコードを眺める

def train(
    self,
    resume_from_checkpoint: Optional[Union[str, bool]] = None,
    trial: Union["optuna.Trial", Dict[str, Any]] = None,
    ignore_keys_for_eval: Optional[List[str]] = None,
    **kwargs,
):

def _inner_training_loop(
    self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None
):
    self.accelerator.free_memory()
    self._train_batch_size = batch_size
    logger.debug(f"Currently training with a batch size of: {self._train_batch_size}")
    # Data loader and number of training steps
    train_dataloader = self.get_train_dataloader()

return DataLoader(
            train_dataset,
            batch_size=self._train_batch_size,
            sampler=train_sampler,
            collate_fn=data_collator,
            drop_last=self.args.dataloader_drop_last,
            num_workers=self.args.dataloader_num_workers,
            pin_memory=self.args.dataloader_pin_memory,
            worker_init_fn=seed_worker,
        )

[
  {
    'input_ids': [0, 0, 0, 0, 0...],
    'labels': [0, 0, 0, 0, 0...],
    'attention_mask': [0, 0, 0, 0, 0...]
  }, 
  {
    'input_ids': [0, 0, 0, 0, 0...],
    'labels': [0, 0, 0, 0, 0...],
    'attention_mask': [0, 0, 0, 0, 0...]
  }
]

 for step, inputs in enumerate(epoch_iterator):

{
  'input_ids': tensor(batch_size, embedding_len),
  'labels': tensor(batch_size, embedding_len),
  'attention_mask': tensor(batch_size, embedding_len)
}

data_collator=DataCollatorForSeq2Seq(
            tokenizer, 
            pad_to_multiple_of=None,
            return_tensors="pt",
            padding=True, 
            label_pad_token_id=tokenizer.pad_token_id,
        ),

class DataLoader(Generic[T_co]):

    def __iter__(self) -> '_BaseDataLoaderIter':
        # When using a single worker the returned iterator should be
        # created everytime to avoid reseting its state
        # However, in the case of a multiple workers iterator
        # the iterator is only created once in the lifetime of the
        # DataLoader object so that workers can be reused
        if self.persistent_workers and self.num_workers > 0:
            if self._iterator is None:
                self._iterator = self._get_iterator()
            else:
                self._iterator._reset(self)
            return self._iterator
        else:
            return self._get_iterator()

return self._get_iterator()

 def _get_iterator(self) -> '_BaseDataLoaderIter':
     ...
     return _SingleProcessDataLoaderIter(self)

class _SingleProcessDataLoaderIter(_BaseDataLoaderIter):
    def __init__(self, loader):

self._dataset_fetcher = _DatasetKind.create_fetcher(
        self._dataset_kind, 
        self._dataset, 
        self._auto_collation, 
        self._collate_fn, 
        self._drop_last)

def _next_data(self):
    index = self._next_index()  # may raise StopIteration
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration

@staticmethod
def create_fetcher(kind, dataset, auto_collation, collate_fn, drop_last):
...
return _utils.fetch._MapDatasetFetcher(dataset, auto_collation, collate_fn, drop_last)

class _MapDatasetFetcher(_BaseDatasetFetcher):
    def fetch(self, possibly_batched_index):
        ...
        return self.collate_fn(data)

def _next_index(self):
    return next(self._sampler_iter)  # may raise StopIteration

@property
def _index_sampler(self):
    # The actual sampler used for generating indices for `_DatasetFetcher`
    # (see _utils/fetch.py) to read data at each time. This would be
    # `.batch_sampler` if in auto-collation mode, and `.sampler` otherwise.
    # We can't change `.sampler` and `.batch_sampler` attributes for BC
    # reasons.
    if self._auto_collation:
        return self.batch_sampler
    else:
        return self.sampler

return DataLoader(
    train_dataset,
    batch_size=self._train_batch_size,
    sampler=train_sampler,
    collate_fn=data_collator,
    drop_last=self.args.dataloader_drop_last,
    num_workers=self.args.dataloader_num_workers,
    pin_memory=self.args.dataloader_pin_memory,
    worker_init_fn=seed_worker,
)

return RandomSampler(self.train_dataset, generator=generator)

class RandomSampler(Sampler[int]):

@dataclass
class DataCollatorForSeq2Seq:

def __call__(self, features, return_tensors=None):

[
  {
    'input_ids': [...],
    'labels': [...],
    'attention_mask': [...]
  },
  {
    'input_ids': [...],
    'labels': [...],
    'attention_mask': [...]
  }
]

  features = self.tokenizer.pad(
            features,
            padding=self.padding,
            max_length=self.max_length,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors=return_tensors,
        )

{
  'input_ids': tensor(batch_size, embedding_len),
  'labels': tensor(batch_size, embedding_len),
  'attention_mask': tensor(batch_size, embedding_len)
}

for step, inputs in enumerate(epoch_iterator):
  ...

{
  'input_ids': tensor(batch_size, embedding_len),
  'labels': tensor(batch_size, embedding_len),
  'attention_mask': tensor(batch_size, embedding_len)
}

for epoch in range(epochs_trained, num_train_epochs):
  ...
  for step, inputs in enumerate(epoch_iterator):
     ...

data = load_dataset("json", data_files=data_path)  
train_val = data["train"].train_test_split(test_size=val_set_size, shuffle=True, seed=42)
train_data = (tokenizerなどの処理)
## listからDatasetに変換
train_data = Dataset.from_list(train_data)

or 

data = load_dataset("json", data_files=data_path)  
train_val = data["train"].train_test_split(test_size=val_set_size, shuffle=True, seed=42)

## mapを使う
train_data = train_val["train"].shuffle().map(generate_and_tokenize_prompt)

from torch.utils.data.dataloader import DataLoader
from torch.utils.data.sampler import RandomSampler
import torch
import transformers
from datasets import load_dataset

data = load_dataset('json', data_files="./sample.json")
train_val = data["train"].train_test_split(
            test_size=1, shuffle=True, seed=42
)
train_data = train_val["train"].shuffle().map(any_function)
# mapを使うとDataSet型のまま
# train_dataが配列の場合はDataSet型にしておく
# train_data = Dataset.from_list(train_data)

sampler = RandomSampler(train_data, generator=torch.Generator())

data_collator=transformers.DataCollatorForSeq2Seq(
    tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
)
data_loader=DataLoader(
    train_data,
    batch_size=4,
    collate_fn=data_collator,
    sampler=sampler,    
)
for steps, input in enumerate(data_loader):    
    print(steps, input['input_ids'].shape, input['labels'].shape)
print('end')