TrainLoader

class lib.training.data.loader.TrainLoader(input_size: int, output_sizes: tuple[int, ...], color_order: T.Literal['bgr', 'rgb'], config: TrainConfig, sampler: None | type[tch_data.RandomSampler | tch_data.DistributedSampler] = None)

Bases: object

Generator for feeding faceswap models with multiple inputs and outputs. Gets the next items from each of the configured loaders and collates them for feeding into a model

Parameters:

  • input_size (int) – The input size to the model

  • output_sizes (tuple[int, ...]) – The output sizes to the model (list as some models have multi-scale outputs)

  • color_order (T.Literal['bgr', 'rgb']) – The color order of the model

  • config (TrainConfig) – The training configuration for feeding the model

  • sampler (None | type[tch_data.RandomSampler | tch_data.DistributedSampler]) – The sampler to use for the data loaders. Default: None (RandomSampler)

Methods Summary

get_loader()

Obtain the dataloaders for each input/output for the model

Methods Documentation

get_loader() → DataLoader

Obtain the dataloaders for each input/output for the model

Return type:

The Training data loaders in side order

__next__() → tuple[list[torch.Tensor], list[torch.Tensor], BatchMeta]

Obtain the next outputs from the loader

Returns:

  • inputs – list of len (num_inputs) tensors of shape(batch_size, H, W, C) inputs for the model

  • targets – List of len (num_outputs) of target images in shape (batch_size, num_inputs, height, width, 3) at all model output sizes as float32 0.0 - 1.0 range

  • meta – The meta information for the batch

Return type:

tuple[list[torch.Tensor], list[torch.Tensor], BatchMeta]

get_loader() → DataLoader

Obtain the dataloaders for each input/output for the model

Return type:

The Training data loaders in side order