New APIs for Transform, Random Transform and SSL, new modules, and transition proposal.

[GabrielBG0]

API Change: New APIs for Transform, Random Transform and SSL, new modules, and transition proposal.

Hello everyone, Gabriel here.

As many of you know, we are gearing up to launch Minerva 1.0. With that there are some changes that have been proposed to fix some of the problems Minerva has currently.

This issue has been written to present some of the changes that have been made and to collect feedback on them. I ask all of you to read this with attention and consideration. If you have problems, questions or suggestions, this is the time to voice them.

TL;DR

Minerva is exiting beta and launching version 1.0, bringing several breaking changes:

• Transforms now accept both data and labels together and return a dictionary instead of a tuple, fixing sync and racing issues in segmentation and contrastive SSL workflows.
• Random Transforms follow the same new API, eliminating the need to sync separate function calls.
• A new SSL interface (_SSLTechnique) standardizes how SSL techniques are implemented across the library.
• Two new modules are being introduced — SSL (for Lightning SSL techniques) and Components (for nn.Module building blocks) — to reduce clutter in the existing Models module.
• Other breaking changes (e.g., variable renames) are also welcome now, as this is the last window for them until version 2.0.
• Workflow change: development moves entirely to public Minerva via forks and PRs, which also fixes the authorship attribution problem caused by the previous private-repo merge process.

Context

As said before, we are preparing to launch Minerva 1.0. Maybe some of you were not aware, but all this time Minerva has been on beta. It was a way for us to evaluate which features are important, what works and what doesn't.

Because of that, some of the implementation decisions that were made have not panned out as intended. We also did not establish a design pattern for some of our features since we did not know what was the best way to do it.

Now the time has arrived to apply what we learn first building minerva to make a better version of it.

Changes

This post will mainly talk about changes in API and in architecture. I will also propose a way to transfer things from our private repo to this one, since it will cease to exist (it will be archived). With that said, let's start with the changes.

Changes on how Transforms work

Transforms on minerva currently follow the pytorch style, applying transforms to one object at a time. This works fine in the pytorch context but since we often work with semantic segmentation there is a need to have synced transforms for both data and label. The solution implemented was less than optimal, having racing problems.

There is also a problem when using contrastive SSL techniques. We needed to create alternative ways for a transform to return multiple views from the same data, and the way we initially implemented transforms did not cover this use case.

With these problems in mind we created a proposal that aims to solve these friction points. The new _Transform API would look like this:

class _Transform: 
    def __init__(self, *args, **kwargs):  
        pass

    def __call__(self, data: dict[str, any], *args, **kwargs) -> dict[str, any]:  
        return {"data": data["data"], "label": data["label"]}

The changes to transform are two fold:

1. It now returns a dictionary
   1. The dictionary always have data and labels keys (even if label is null)
   2. If you need multiple views for the data, just add to the dict!
   3. Your model will now access keys from this dict instead of positions on a tuple.
2. Transforms now accept both data and labels at the same time.
   1. If your transform may be only needed to data or label separately you can add a flag to control this behavior.
   2. Transforms now can be much more easily synced for data and labels since they are processed together.
   3. No racing issues.

These changes to transforms will also impact our datasets and readers, as they will also return dictionaries. There will also not be a need to pass an individual transform pipeline for data and others for labels.

Changes to Random Transforms

As stated before, syncing random transforms for data and labels has been a problem, especially so in distributed environments. The changes into the transform API will allow us to better manage random transforms. This is the new proposed API for random transform:

class _RandomTransform():
    def __init__(self, parameters, p, seed):
        """Apply random transform to provided data


        you can set parameters for the transform,
        the probability of applying the transform and the seed for reproducibility
        """
        pass

    def __call__(self, data: dict[str, any]) -> dict[str, any]:
        """ Apply the transformation to the provided data
        This function will generate the random parameters for the transformation and apply it to the data.
        """
        return {"data": data["data"], "label": data["label"]}

Contrastive transforms and context transforms will follow this same API, only adapting to their specific needs.

Since the call function has both data and label there will not be a need to sync the transform between multiple function calls.

New SSL Interface

In the current state Minerva does not have an SSL technique interface that standardizes our ssl implementations. This made the techniques implementations to not have a clear pattern and thus, making them more difficult to use. Considering this issue here is the proposed interface for ssl techniques:

_SSLTechnique — Abstract Base Class

_SSLTechnique is an abstract base class that extends PyTorch Lightning's LightningModule. It serves as the foundation for all SSL (Self-Supervised Learning) technique implementations.

Constructor Parameters

The class is initialized with the following parameters:

• backbone (nn.Module) — The feature-extraction network shared across all components of the technique.
• learning_rate (float, default 1e-3) — The base learning rate passed to the default Adam optimizer.
• train_metrics (Dict[str, Metric], default None) — A dictionary of Torchmetrics to log during training. Requires _predictions_step to be overridden.
• val_metrics (Dict[str, Metric], default None) — Torchmetrics to log during validation.

Abstract Methods

These methods have no default implementation and must be overridden by every subclass:

• forward(x) → Any — Defines the full forward pass. The return type is technique-specific and may be a tensor, a tuple, or another structure.
• _loss_step(batch) → Tensor — Computes the scalar loss for a single batch. This method is invoked internally by the training, validation, and test step methods.
• training_step(batch, idx) — Implements the Lightning training step following technique-specific logic.
• validation_step(batch, idx) — Implements the Lightning validation step following technique-specific logic.

Overridable Hooks

These methods have sensible defaults and can optionally be overridden to customize behavior:

• _predictions_step(batch) — Returns a (y_hat, y) tuple for metric evaluation. Defaults to None; should only be overridden when the technique includes a supervised head.
• technique_transforms() — Returns a TransformPipeline with technique-specific data augmentations. Defaults to None (no augmentations).
• technique_callbacks() — Returns a list of Lightning callbacks with technique-specific behavior. Defaults to None (no callbacks).
• technique_collate_fn() — Returns a custom batch collation function for DataLoaders. Useful when the technique requires non-standard batch formatting. Defaults to None (standard collation).
• default_train_strategy() — Returns a string indicating the recommended training strategy. Useful for techniques that require a specific distributed strategy. Defaults to "auto" (no specific recommendation).
• configure_optimizers() — Returns the optimizer configuration. Defaults to Adam. Override to substitute SGD, LARS, or to attach a learning rate scheduler.

Concrete Methods

These methods are fully implemented by the base class and available to all subclasses without any changes:

• test_step(batch, idx) — Runs _loss_step and any registered test metrics, then logs test_loss.
• _compute_metrics(y_hat, y, phase) — Evaluates all metrics registered for the given phase and returns a prefixed logging dictionary.
• get_representations(x) — Returns raw backbone embeddings. Intended for downstream evaluation tasks without requiring direct access to the model's internals.

New Modules

We also identified some problems with how our models, ssl techniques and nn modules were organised. Currently, we have a single model module that hosts all models (lightning and nn.modules) and ssl techniques. This made the model module very crowded and with a lot of sub divisions. To address this issue we decided to make two new root modules:

1. SSL: that hosts the lightning ssl technique implementations
2. Components (im still not too sure on the name, suggestions are welcomed): will house the nn.Modules implementations that are used to build the lightning models we have.

So, the new structure on minerva will be:

• Models: house the lightning implementations for our models
• SSL: house the lightning implementations for our ssl techniques
• Components(?): house the pieces (or blocks) used to build the ssl or model implementations in lightning. These are the classes that inherit from nn.Module.

So, what changes? A few things:

1. If you have a model implementation file that houses both nn.Module implementations and LightningModel implementations this will be split into two files. (one in the Models module and the other in the Components module)
2. If you have a SSL technique implementation that houses both nn.Module implementations and LightningModule implementations they will also be split in two files. (one in the SSL module and the other in the Components module)
3. We (hope) that with this separation code reuse will be more effective, allowing implementations of lightning models that use the same type of backbone or prediction head to be able to share the same component class. (currently we have 3 or 4 ViT implementations, wich can make maintenance very difficult)

Other changes

Some other braking changes will be implemented, like changing the name of the fc variable on SimpleSupervisedModel changing to predictions_head. If you have a change of this nature (changing variable names for a better one, or something like that) this is the time to do it. I only ask you to open an issue on github to register the change.

This will be the last time braking changes will be accepted for a while (until version 2.0) so the time to make them is now. If you have a problem with how the library works currently and the changes addressed in this post won't be solved then this is the time to contact us and tell us your problem. We can’t guarantee that we will be able to solve your problem but this is the time for us to evaluate what we want to include in version 1.0.

While version 1.0 matures other changes may be brought forward, if that happens we will contact you for feedback again.

Transition proposal and new development practices

As you can see, there are a lot of changes to be made. There was also brought to the maintenance team a problem with authorship recognition. As you know, development has been made in our private branch and merged to the public one by Otavio. This has created a problem where every contribution to the public repo was being made in his name, and the name of the original authors were being lost in the process.

We can solve these two problems with one solution. The original author from each technique will adapt the code to the new standard and make any changes necessary. They will then open a PR into public minerva with the changes. This way their contributions will be registered as theirs. We won't be able to bring their past contributions to their name but from here forward they will appear as contributors to the project. If your code is not yet in the public minerva I advise you to contact your supervisor and check if you can make it public.

From now on the work process will be a little different. Now contributions will only be accepted into public minerva. To do that you will need to make a fork of minerva and implement your changes there. When you are done you can submit your changes through a PR. If the work you are developing is confidential and should not be publicly available until publication we strongly suggest that you private your fork. If you have any questions about workflow you can contact me.

Once more I ask you to voice your opinion on these changes and any other you may consider beneficial for Minerva.

Thank you for the time and attention,

Gabriel G & Minerva Core Team

[filipeas]

Understood. About ViT versions, this will impact me more than this Transforms, because a follow the default Pipeline test (suggested by Otavio).

[otavioon]

Very nice! I just raise some qusestions:

---

technique_callbacks() — Returns a list of Lightning callbacks with technique-specific behavior. Defaults to None (no callbacks).
technique_collate_fn() — Returns a custom batch collation function for DataLoaders. Useful when the technique requires non-standard batch formatting. Defaults to None (standard collation).
default_train_strategy() — Returns a string indicating the recommended training strategy. Useful for techniques that require a specific distributed strategy. Defaults to "auto" (no specific recommendation).

This would be called before instantiating trainer/datamodule and would be passed to them, right?

---

get_representations(x) — Returns raw backbone embeddings. Intended for downstream evaluation tasks without requiring direct access to the model's internals.

I would use get_embeddings(x) instead of representations. It's more common.

---

y_hat refer to logits, right?

[eborin]

Gabriel, thanks for the discussion. Here are my thoughts.

CHANGES

Changes on how Transforms work

Transferring the responsibility of producing multiple views to the Transform class would require this class to be updated every time a new SSL method that requires a new kind of view is implemented. It seems more natural to me to transfer this responsibility to a Dataset class that encapsulates the original dataset and produces the views needed by the SSL method (see below).

A key aspect that should be discussed is the type of the object returned by getitem, as some of the SSL methods have assumptions about these objects and Pytorch datasets do not specify their types.
For example, TriBYOL expects three views for each sample on the batch, while BYOL expects two views, and LFR expects only the sample itself. The implication is that given an SSL technique, the dataset must be adapted to work with this technique (for example, employ transformations that generate two or three views of each item so the getitem method returns these views, instead of the data item itself).

In this context, it seems more appropriate to create dataset classes specialized for each SSL technique.
For example, we could have a TriBYOL_Dataset class that encapsulates Pytorch datasets and apply transformations to produce tuples with three views to be returned by the getitem method. Something like this:

from torch.utils.data import Dataset
class TriBYOL_Dataset(Dataset):

       def __init__(
               self,
               dataset,
               transforms: Optional[Union[_Transform, List[_Transform]]] = None,
           ):
              self.dataset = dataset
              if transforms is None:             
                 self.transforms = ....  # Default TriBYOL transforms
             else:
                 self.transforms = transforms
           

       def __len__(self) -> int:
              raise NotImplementedError

       def __getitem__(self, idx: int) -> Union[Any, Tuple[Any, ...]]:
              item = self.dataset[idx]
              return (gen_view(item), gen_view(item), gen_view(item))

       def gen_view(item):
              x = item
              for t in transforms: x = t(x)
              return x

In this way, the training workflow would follow a consistent, three-line pattern regardless of the SSL technique:

• Example: TriBYOL Training

tribyol_model = TriBYOL(backbone)
tribyol_ds = TriBYOL_Dataset(dataset)
pl.Train(tribyol_model, tribyol_ds)

• Example: BYOL Training -- training the same backbone on the same dataset with another SSL technique simply:

byol_model = BYOL(backbone)
byol_ds = BYOL_Dataset(dataset)
pl.Train(byol_model, byol_ds)

Changes to Random Transforms

It seems to me that receiving a raw seed is problematic. Which generator exactly would be set with this seed? Wouldn't setting a global generator risk overwriting the settings performed by other parts of the code?

Instead of a seed, how about we allow the user to pass an explicit random number generator object? In this way, the user has much more control over the reproducibility process. Furthermore, we reduce the chance of other code breaking its own assumptions about random number generation, which often happens when multiple seed initializations occur on the same global generator.

Regarding the Transform returning a dictionary, I think we should keep the code compatible with standard PyTorch transforms. This way, we can interchangeably use their transforms without issues.

New SSL Interface

I agree we should have an interface for SSL techniques. However, I would like to understand better the motivation behind defining train_metrics and val_metrics on the class constructor.
Also, why specify the learning_rate on the constructor? Shouldn’t it be passed on optimizer_kwargs, together with other optimizer arguments?

Abstract Methods

What is the purpose of the forward() method on the SSL class? It should run the input data through the entire pretext model (backbone + projection head) and produce the pretext labels, run just the backbone and return the features?

Why require the implementation of the _loss_step method?

Overridable Hooks

Why implement the _prediction_step(batch) method? Is it required to learn representations?

technique_transforms(): Some SSL techniques can be applied to different data domains (e.g., audio, photos, seismic data, ACC/GYR time series, …). The transformations applied on different data domain may vary, and defining a default on the SSL class imply that it is being specialized to some domain. Instead, I think it would be more natural to transfer this responsibility to a Dataset class that encapsulates the original dataset and defines the transforms to be applied on the data (as discussed above).

What is the reason to include technique_callbacks in the API?

What is the reason to include technique_collate_fn in the API? Any use cases that would make this essential?

Concrete Methods

Why implement the test_step? Are we assuming the SSL methods will be executed on test data?

_compute_metrics(): It seems to me that metrics used on SSL method are tightly coupled with the pretext task and very particular to each SSL method. In this context, does it make sense to try to standardize the metrics used by these methods?

New Modules

I think Components is too generic to represent model basic blocks.
How about using nn (Neural Networks), as pytorch uses?

Regarding splitting SSL technique implementation that houses both nn.Module implementations and LightningModule implementations into two files, if the nn.Module is an private object, not exported on the API, I believe it should be kept in the same file as the SSL technique.

Other changes

why predictions_head? Shouldn’t it be prediction_head?

[GabrielBG0]

@otavioon

This would be called before instantiating trainer/datamodule and would be passed to them, right?

Yes, they would.

I would use get_embeddings(x) instead of representations. It's more common.

I agree, thanks for the suggestion

y_hat refer to logits, right?

Yes

[GabrielBG0]

@eborin I'm finishing an in-depth description of the transform and dataset changes. In the meantime, I'll address your questions regarding the SSL interface.

First, I want to clarify that the SSL interface was designed by analyzing the current SSL classes and attempting to accommodate their existing features as consistently as possible.

I would like to understand better the motivation behind defining train_metrics and val_metrics on the class constructor.

This decision comes from the current implementation of TFC as it defines train_metrics, val_metrics, and test_metrics. Since I wasn't sure whether this is specific to TFC or broadly applicable to other SSL techniques, I included it as a general option.

why specify the learning_rate on the constructor? Shouldn’t it be passed on optimizer_kwargs, together with other optimizer arguments?

We could include it in optimizer_kwargs, but I consider the learning rate important enough to be exposed as a top-level parameter. This is also consistent with how we define our models, so it improves uniformity across the codebase.

What is the purpose of the forward() method on the SSL class? It should run the input data through the entire pretext model (backbone + projection head) and produce the pretext labels, run just the backbone and return the features?

The forward method is intended to execute the full pipeline: backbone → optional intermediate processing → projection head → return projected representations.

Why require the implementation of the _loss_step method?

_loss_step consumes the output of forward and computes the loss. Since this logic is shared between training_step and validation_step, it promotes code reuse. It also provides a clear hook for any additional preprocessing required before loss computation.

Why implement the _prediction_step(batch) method? Is it required to learn representations?

It is not required. However, if a technique includes a prediction head, this provides a standardized way to expose predictions. As defined in the spec, it is optional and intended as a convenience feature.

technique_transforms(): Some SSL techniques can be applied to different data domains (e.g., audio, photos, seismic data, ACC/GYR time series, …).

Agreed. Based on our discussion, this is likely not the appropriate place to define transformations. Creating separate classes per technique and per data domain is also not ideal. This will be better addressed alongside the dataset and transform refactor.

What is the reason to include technique_callbacks in the API?

This allows the technique author to specify recommended training configurations (e.g., callbacks), making best practices more explicit to the user.

What is the reason to include technique_collate_fn in the API? Any use cases that would make this essential?

This was suggested by @otavioon. I am not fully familiar with the use cases yet.

Why implement the test_step? Are we assuming the SSL methods will be executed on test data?

_compute_metrics(): It seems to me that metrics used on SSL method are tightly coupled with the pretext task and very particular to each SSL method. In this context, does it make sense to try to standardize the metrics used by these methods?

Both points were inherited from the TFC implementation. Since I was unsure of their broader necessity, I included them provisionally. If they are not generally applicable, we can remove them.

I think Components is too generic to represent model basic blocks.

Agreed.

How about using nn (Neural Networks), as pytorch uses?

On pytorch the model "blocks" are under nn.module not nn. I find "Neural Networks" slightly misleading here, since we are exposing building blocks rather than full networks. "ModelComponents" might be a better alternative, but this is open for discussion.

Regarding splitting SSL technique implementation that houses both nn.Module implementations and LightningModule implementations into two files, if the nn.Module is an private object, not exported on the API, I believe it should be kept in the same file as the SSL technique.

First, to clarify: Python does not enforce true privacy. The leading underscore (_) is only a convention indicating internal use, so every class/function (etc...) is exported in the API.

Second, I don't see why we would separate nn.module classes from our Lightning Module classes but not do it for our SSL techniques.

why predictions_head? Shouldn’t it be prediction_head?

This is a typo, prediction_head is correct.

[GabrielBG0]

@eborin sorry for the late reply, here is mine and @otavioon's proposition for using dict on dataset and transforms.

Change from tuples to dictionaries

In this post, we argue why we need the changes we proposed, more specifically, the change from datasets and transforms using/returning tuples to using/returning dictionaries.

Why do we need to change from tuples to dictionaries?

In the current framework, during development and use we started to encounter some problems with how our data pipeline works.
The main problem is that it was very hard to sync transforms between data and label.

Transforms may need more than one element as an input

This is evident from two scenarios:

1. In contexts like semantic segmentation, where you have a dense prediction, you need to apply the same transform to both the image and the label.
2. Some transforms are conditional and depend on the label, for example, in semantic segmentation, you might want to apply a random crop that tries to find a crop where the majority class does not occupy more than X% of the crop, this is a common technique to avoid having too many crops with only background class.

The natural solution to this problem is to design the transforms to be able to take more than one element as an input, and to return more than one element as an output, this way we can ensure that the transforms are applied in a consistent way to both the data and the label and that we can have transforms that are conditional on the label.

Using dictionaries is the natural way to implement this

Dictionaries are the natural way to implement this. With them, debugging is much easier and development is more consistent, since we do not need to worry about the order of the elements in the tuple, and we can easily add new elements to the dictionary without breaking existing code.

When using dictionaries, the user can specify the keys that they want to be transformed (for cases where the label is not "transformable" like in detection, or transforms that only make sense for the data and not for the label, like color jitter or gaussian blur). Any keys not listed in keys_to_transform are passed through unchanged, so the output dictionary always contains all original keys.

The __init__ of the transform would then take a list of keys that it should transform, and the __call__ method would apply the transform only to those keys, like this:

class CropTransform(_Transforms):
    def __init__(self, keys_to_transform: List[str], crop_size: Tuple[int, int]):
        self.keys_to_transform = keys_to_transform # <- user specifies which keys to transform
        self.crop_size = crop_size

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        # apply the crop transform to the specified keys and return a dictionary with the transformed data
        # keys not in keys_to_transform are left unchanged
        for key in self.keys_to_transform:
            data[key] = self.crop(data[key], self.crop_size) # <- simplification for the sake of the example
        return data

This would also facilitate the implementation and use of multi-view transforms. When a multi-view transform is applied, the user can specify new keys for the new views, and the transform would return a dictionary with the new views as new keys, like this:

class MultiViewTransform(_Transforms):
    def __init__(self, keys_to_transform: List[str], new_keys: List[str]):
        self.keys_to_transform = keys_to_transform # <- user specifies which keys to transform
        self.new_keys = new_keys # <- user specifies the new keys for the new views

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        # apply the transform to the specified keys and return a dictionary with the new views as new keys
        for key in self.keys_to_transform:
            new_views = self.transform(data[key]) # <- simplification for the sake of the example
            for i, new_key in enumerate(self.new_keys):
                data[new_key] = new_views[i] # <- add the new view to the dictionary with the new key
        return data

It would also be possible to have contraction transforms, where the transform takes multiple keys as input and returns a single key.

class ContractionTransform(_Transforms):
    def __init__(self, keys_to_transform: List[str], new_key: str):
        self.keys_to_transform = keys_to_transform # <- user specifies which keys to transform
        self.new_key = new_key # <- user specifies the new key for the contracted view

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        # apply the transform to the specified keys and return a dictionary with the contracted view as a new key
        inputs = [data[key] for key in self.keys_to_transform] # <- get the inputs for the transform
        contracted_view = self.transform(inputs) # <- simplification for the sake of the example
        data[self.new_key] = contracted_view # <- add the contracted view to the dictionary with the new key
        return data

As for random transforms, the implementation and synchronization would be trivial:

class RandomCropTransform(_Transforms):
    def __init__(self, keys_to_transform: List[str], crop_size: Tuple[int, int]):
        self.keys_to_transform = keys_to_transform # <- user specifies which keys to transform
        self.crop_size = crop_size

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        # apply the random crop transform to the specified keys and return a dictionary with the transformed data
        # we can use the same random crop parameters for all the keys to ensure synchronization
        crop_params = self.get_random_crop_params(data[self.keys_to_transform[0]], self.crop_size) # <- get the random crop parameters from the first key (simplification for the sake of the example)
        for key in self.keys_to_transform:
            data[key] = self.crop(data[key], crop_params) # <- apply the same crop parameters to all the specified keys
        return data

How to ensure compatibility with external code?

It is true that the "normal" way to implement datasets and transforms in PyTorch is to use tuples, and that using dictionaries might break compatibility with external code that expects tuples.

However, we can easily ensure compatibility by providing wrappers that convert dictionaries to tuples and vice versa. This approach is also compatible with torchvision.transforms.Compose, since wrapping a legacy pipeline with TupleToDictTransformWrapper allows it to be used transparently inside our dictionary-based _TransformsPipeline. This way, users can still use external code that expects tuples without any modification, while we can still have the benefits of using dictionaries internally.

class TupleToDictTransformWrapper(_Transforms):
    def __init__(self, keys: List[str], transform):
        self.keys = keys # <- user specifies the keys to convert to tuples
        self.transform = transform # <- the transform to apply to the tuples

    def __call__(self, data: Union[Tuple[Any, ...], Any]) -> Dict[str, Any]:
        # convert the specified keys to tuples, apply the transform, and convert back to dictionaries

        if isinstance(data, tuple):
            tuple_data = tuple(data[key] for key in self.keys)
            transformed_data = self.transform(tuple_data)
        else:
            transformed_data = self.transform(data) # <- if the input is already a dictionary, we can apply the transform directly

        return {self.keys[0]: transformed_data} if not isinstance(transformed_data, tuple) else {key: transformed_data[i] for i, key in enumerate(self.keys)} # <- convert back to dictionary with the same keys

class DictToTupleTransformWrapper(_Transforms):
    def __init__(self, keys: List[str], transform):
        self.keys = keys # <- user specifies the keys to convert to dictionaries
        self.transform = transform # <- the transform to apply to the dictionaries

    def __call__(self, data: Tuple[Any, ...]) -> Tuple[Any, ...]:
        # convert the tuple to a dictionary based on the specified keys
        data_dict = {key: data[i] for i, key in enumerate(self.keys)}
        transformed_dict = self.transform(data_dict)
        return tuple(transformed_dict[key] for key in self.keys) # <- convert back to tuple

Example of how to use the wrappers:

# Example 1: Using a PyTorch transform that expects a single element with our new dictionary-based implementation

color_jitter = torchvision.transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1) # <- a PyTorch transform that expects a single element (the image)
color_jitter_wrapper = TupleToDictTransformWrapper(keys=["image"], transform=color_jitter) # <- wrap the PyTorch transform to convert the dictionary to a tuple and back
transforms = _TransformsPipeline([
    CropTransform(keys_to_transform=["image", "label"], crop_size=(256, 256)),
    RandomCropTransform(keys_to_transform=["image", "label"], crop_size=(128, 128)),
    color_jitter_wrapper, # <- use the wrapped PyTorch transform in our transform pipeline
])

# Example 2: using our new dictionary-based implementation with an external dataset that would pass a single element to the transform

gaussian_blur = GaussianBlurTransform(keys=["image"], kernel_size=5, sigma=(0.1, 2.0)) # <- our new dictionary-based implementation of a Gaussian blur transform that expects a dictionary with a single key "image"
gaussian_blur_wrapper = TupleToDictTransformWrapper(keys=["image"], transform=gaussian_blur) # <- wrap our new transform to convert the tuple to a dictionary and back

external_dataset = ExtDataset(path="data", transform=gaussian_blur_wrapper) # <- an external dataset that sends a single element (the image) to the transform, we wrap our new transform to convert the tuple to a dictionary and back, so we can use it with the external dataset without any modification

This mode of operation would allow us to have the best of both worlds: we can use dictionaries internally for better development and debugging experience, while still being able to use external code that expects tuples without any modification.

On the transform documentation, the writer would specify the type of the transform, whether it is a n→n transform, a multi-view transform, or a contraction transform, and the expected keys in the input and output dictionaries. This way, users can easily understand how to use the transform and what to expect as input and output.

For the multi-view and contraction transforms, it would be possible to use string matching to specify the new keys for the new views, for example, if the input dictionary has keys "image" and "label", and we want to create two new views for the image, we could specify the new keys as "image_view*", and the transform would automatically create new keys "image_view1" and "image_view2" for the new views. This would make it easier for users to specify the new keys without having to worry about the exact names of the keys.
When receiving a dictionary with multi-view keys, you could also specify the transform keys as "image_view*", and the transform would apply to all the keys that match the pattern "image_view*". This would make it easier for users to apply transforms to all the views without having to specify each key individually.

Impacts on the dataset

The main impact on the dataset would be that the dataset would need to return a dictionary instead of a tuple. This would require some changes in the implementation of the dataset, but it would also allow for more flexibility in the data pipeline, as we can easily add new elements to the dictionary without breaking existing code.

As proposed changes we have:

1. Datasets now return a dictionary instead of a tuple, where the keys are specified by the user and can be anything (e.g., "data", "label", "metadata", etc.).
2. Datasets now take an OrderedDict of readers instead of a list of readers, where the keys are specified by the user and can be anything (e.g., "data", "label", "metadata", etc.). The dataset would use the keys to read the data and return a dictionary with the same keys and the corresponding data as values.
3. Instead of a list of _Transforms, one for each reader, the dataset would take a single _Transforms that would be applied to the dictionary returned by the readers. The _Transforms would be responsible for applying the appropriate transforms to the appropriate keys in the dictionary, based on the user specifications.

Example of the definition of the new dataset:

class MinervaDataset(torch.utils.data.Dataset):
    def __init__(self, readers: OrderedDict[str, Reader], transforms: _Transforms):
        self.readers = readers # <- user specifies the readers with their corresponding keys
        self.transforms = transforms # <- the transforms to apply to the dictionary returned by the readers (ether a single transform or a transform pipeline)

    def __len__(self):
        # return the length of the dataset based on one of the readers (assuming all readers have the same length)
        return len(next(iter(self.readers.values())))

    def __getitem__(self, idx) -> OrderedDict[str, Any]:
        # read the data from the readers and return a dictionary with the data
        data = OrderedDict()
        for key, reader in self.readers.items():
            data[key] = reader[idx] # <- read the data for each key using the corresponding reader
        transformed_data = self.transforms(data) # <- apply the transforms to the dictionary
        return transformed_data # <- return the transformed dictionary

This new implementation of the dataset would allow for more flexibility in the data pipeline, as we can easily add new elements to the dictionary without breaking existing code, and we can also have more complex transforms that can take multiple keys as input and return multiple keys as output.
We would also be able to reuse the same dataset for different tasks by simply changing the transforms, readers and the keys in the dictionary, without having to change the implementation of the dataset itself, or creating new datasets for each task. This would make it easier for users to experiment with different tasks and data pipelines without having to worry about the implementation of the dataset.

Some examples of how to use the new dataset:

# Example 1: A simple dataset for image segmentation, where we have an image reader and a label reader, and we want to apply the same transforms to both the image and the label.
image_reader = ImageReader(...)
label_reader = LabelReader(...)
readers = OrderedDict({"image": image_reader, "label": label_reader})
transforms = _TransformsPipeline([
    CropTransform(keys_to_transform=["image", "label"], crop_size=(256, 256)),
    RandomCropTransform(keys_to_transform=["image", "label"], crop_size=(128, 128)),
    ColorJitterTransform(keys_to_transform=["image"], brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
])
dataset = MinervaDataset(readers=readers, transforms=transforms)

# Example 2: A dataset for multi-view learning, where we have a single image reader, and we want to create two new views for the image using a multi-view transform.
image_reader = ImageReader(...)
readers = OrderedDict({"image": image_reader})
transforms = _TransformsPipeline([
    MultiViewTransform(keys_to_transform=["image"], new_keys=["image_view1", "image_view2"]),
    ColorJitterTransform(keys_to_transform=["image_view1"], brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    GaussianBlurTransform(keys_to_transform=["image_view2"], kernel_size=5, sigma=(0.1, 2.0)),
])
dataset = MinervaDataset(readers=readers, transforms=transforms)

# Example 3: A dataset for a task where we have two image readers and two label readers, and we want to create a new view that is a contraction of the two images and the two labels using a contraction transform.
image_reader1 = ImageReader(...)
image_reader2 = ImageReader(...)
label_reader1 = LabelReader(...)
label_reader2 = LabelReader(...)

readers = OrderedDict({"image1": image_reader1, "image2": image_reader2, "label1": label_reader1, "label2": label_reader2})
transforms = _TransformsPipeline([
    ContractionTransform(keys_to_transform=["image1", "image2"], new_key="contracted_image"),
    ContractionTransform(keys_to_transform=["label1", "label2"], new_key="contracted_label"),
])

dataset = MinervaDataset(readers=readers, transforms=transforms)

# Example 4: SAM needs binary masks for labels and we have a data reader and label reader that return the data and the label in a dictionary with keys "data" and "label", we can use a multi-view transform to create a new view for the binary masks. SAM also needs some additional metadata for the labels, we can also add that to the dictionary using a metadata reader.

data_reader = DataReader(...)
label_reader = LabelReader(...)
metadata_reader = MetadataReader(...)

readers = OrderedDict({"data": data_reader, "label": label_reader, "metadata": metadata_reader})
transforms = _TransformsPipeline([
    BinaryMaskExtractor(keys_to_transform=["label"], num_classes=10, new_keys="binary_mask_*"),
])

dataset = MinervaDataset(readers=readers, transforms=transforms)

# note that the returned dictionary from the dataset would have the following keys: "data", "label", "metadata", "binary_mask_1", "binary_mask_2", ... so the model keeps the original label and the new binary masks as separate keys in the dictionary.

Compatibility with external code

In the same way as with the transforms, we can provide a wrapper for the dataset that converts the dictionaries to tuples and vice versa, to ensure compatibility with external code that expects tuples.

class ExternalDatasetWrapper:
    def __init__(self, dataset, keys: List[str]):
        self.dataset = dataset
        self.keys = keys

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx) -> Tuple[Any, ...]:
        data_tuple = self.dataset[idx] # <- get the data as a dictionary from the dataset
        return tuple(data_tuple[key] for key in self.keys) # <- convert the dictionary to a tuple based on the specified keys

This wrapper would allow users to use external code that expects tuples without any modification, while still being able to use the new dataset implementation that returns dictionaries internally. Users can simply wrap the new dataset with this wrapper when they need to use it with external code that expects tuples.

An example of how to use the wrapper:

# Example: Using the new dataset with external code that expects tuples
cifar10_dataset = CIFAR10(root="data", train=True, download=True) # <- an external dataset that returns x, y tuples

dataset = ExternalDatasetWrapper(dataset=cifar10_dataset, keys=["image", "label"])

Conclusion

With dictionaries, we don't need to create specialized datasets for each task. We can simply change the readers and the transforms to adapt to different tasks, and we can also have more complex transforms that can take multiple keys as input and return multiple keys as output. We also accommodate the need for additional metadata in a clean and flexible way depending on the model's needs.

Impacts on models

The main impact on models would be that they would need to unpack a dictionary instead of a tuple as input. The rest of the model implementation would remain mostly unchanged, as the model would still be able to access the data and the labels using the keys in the dictionary.

For legacy models, two aspects require attention:

• Forward pass: Models that expect positional tuple arguments will need to unpack the dictionary before calling the model. A wrapper can handle this transparently.
• Lifecycle methods: training_step, val_step, and predict_step in Lightning-based models also receive the batch from the dataloader, so they equally need to unpack a dictionary rather than a tuple. A lifecycle wrapper can be provided for legacy models to handle this without any changes to the model itself.

For models outside of the framework, we can provide a wrapper that converts the dictionaries to tuples, to ensure compatibility with external code that expects tuples.

class ExternalModelDatasetWrapper(torch.utils.data.Dataset):
    def __init__(self, dataset, keys: List[str]):
        self.dataset = dataset # <- dataset that returns dictionaries with the data and labels
        self.keys = keys

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx) -> Tuple[Any, ...]:
        data_dict = self.dataset[idx] # <- get the data as a dictionary from the dataset
        return tuple(data_dict[key] for key in self.keys) # <- convert the dictionary to a tuple based on the specified keys

This wrapper would allow users to use external code that expects tuples without any modification, while still being able to use the new dataset implementation that returns dictionaries internally. Users can simply wrap the new dataset with this wrapper when they need to use it with external code that expects tuples.

An example of how to use the wrapper:

# Example: Using our new dataset with an external model that expects tuples

readers = OrderedDict({"data": DataReader(...), "label": LabelReader(...)})
transforms = _TransformsPipeline([
    CropTransform(keys_to_transform=["data", "label"], crop_size=(256, 256)),
    RandomCropTransform(keys_to_transform=["data", "label"], crop_size=(128, 128)),
    ColorJitterTransform(keys_to_transform=["data"], brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
])

dataset = MinervaDataset(readers=readers, transforms=transforms)
wrapped_dataset = ExternalModelDatasetWrapper(dataset=dataset, keys=["data", "label"])

model = ExternalModel(...) # <- an external model that expects x, y tuples

train(model, wrapped_dataset, ...)

How do I know what keys to use in the model?

This would be specified in the documentation of the model, where the expected keys in the input dictionary would be clearly stated. For example, if the model expects an input dictionary with keys "data" and "label", the documentation would specify that the model expects a dictionary with those keys, and that the "data" key should contain the input data and the "label" key should contain the corresponding labels. This way, users can easily understand what keys to use in the model and how to prepare their dataset accordingly.

Techniques that require multiple views can specify a string pattern for the expected keys, for example, if the model expects multiple views of the data with keys "data_view1", "data_view2", etc., the documentation would specify that the model expects keys that match the pattern "data_view*", and that the model will use all the keys that match this pattern as input data. This way, users can easily understand how to prepare their dataset for multi-view learning and what keys to use in the model.

Wrappers

All wrappers described in this document, for both datasets and transforms, will be provided by default as part of the framework. Users will not need to implement them themselves.

These wrappers cover two directions:

• Inward: wrapping external datasets and transforms (tuple-based) so they can be used inside our dictionary-based pipeline.
• Outward (inverse): wrapping our datasets so that external PyTorch code or third-party frameworks that expect tuples can consume them without any modification.

The inverse wrapper is particularly useful when integrating a framework model into an external codebase. For example:

# Example: Exposing a framework dataset to an external PyTorch training loop

readers = OrderedDict({"data": DataReader(...), "label": LabelReader(...)})
transforms = _TransformsPipeline([
    CropTransform(keys_to_transform=["data", "label"], crop_size=(256, 256)),
])

dataset = MinervaDataset(readers=readers, transforms=transforms)

# Wrap the dataset so external code receives (data, label) tuples instead of a dictionary
external_ready_dataset = ExternalDatasetWrapper(dataset=dataset, keys=["data", "label"])

# Can now be passed directly to any standard PyTorch training loop
dataloader = torch.utils.data.DataLoader(external_ready_dataset, batch_size=32)

The wrappers are simple to use and provide a seamless path for users to integrate the new dictionary-based pipeline with their existing code in either direction.

Final Conclusion

The transition to a dictionary-based pipeline fundamentally shifts how we handle data preparation. Using this method, we rarely need to create specialized dataset classes (like a TriBYOLDataset, SimCLRDataset, etc.), nor do we even need to write new dataset implementations for day-to-day tasks.

Because the MinervaDataset class already allows for several complex constructions in a highly parametrizable way, inheriting from MinervaDataset will be reserved exclusively for edge cases (like filtering the dataset to have a percentage x from a class). Instead of writing boilerplate dataset code, users can define their data flow purely through the composition of readers and transforms.

This makes the creation of datasets for complex architectures, such as contrastive learning, completely trivial. Contrastive learning methods rely heavily on generating multiple augmented views of the same input, which historically required hardcoding custom __getitem__ logic. Now, we can achieve this simply by leveraging the MultiViewTransform and targeting the new dictionary keys.

Example: Creating a SimCLR Dataset

SimCLR requires two different augmented views of the same image to serve as positive pairs. Instead of a custom dataset class, we can just instantiate MinervaDataset with a pipeline that splits the image into two keys and applies stochastic augmentations to each independently.

# SimCLR Dataset Construction (simplified)
image_reader = ImageReader(path="data/images")
readers = OrderedDict({"image": image_reader})

transforms = _TransformsPipeline([
    # Split the original image into two separate views
    MultiViewTransform(keys_to_transform=["image"], new_keys=["view1", "view2"]),
    
    # Apply standard SimCLR augmentations to the views
    RandomResizedCropTransform(keys_to_transform=["view1", "view2"], size=(224, 224)),
    ColorJitterTransform(keys_to_transform=["view1", "view2"], brightness=0.8, contrast=0.8, saturation=0.8, hue=0.2),
    GaussianBlurTransform(keys_to_transform=["view1"], kernel_size=23, sigma=(0.1, 2.0)), 
    # Notice we can easily apply blur only to view1, or use different parameters for each
])

simclr_dataset = MinervaDataset(readers=readers, transforms=transforms)

# The model's forward pass will receive a dictionary containing "view1" and "view2"

Example: Creating a TriBYOL Dataset

TriBYOL extends the multi-view concept by utilizing three distinct views (often a multi-crop strategy with different resolutions or distinct augmentation pipelines). With dictionaries, scaling to three views requires almost zero additional effort.

# TriBYOL Dataset Construction (simplified)
image_reader = ImageReader(path="data/images")
readers = OrderedDict({"image": image_reader})

transforms = _TransformsPipeline([
    # Split the original image into three separate views
    MultiViewTransform(keys_to_transform=["image"], new_keys=["view1", "view2", "view3"]),
    
    # Apply baseline augmentations
    RandomResizedCropTransform(keys_to_transform=["view1", "view2"], size=(224, 224)),
    RandomResizedCropTransform(keys_to_transform=["view3"], size=(96, 96)), # A local crop view
    
    # Apply asymmetric augmentations easily by targeting specific keys
    ColorJitterTransform(keys_to_transform=["view1", "view2", "view3"], brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1),
    GaussianBlurTransform(keys_to_transform=["view2", "view3"], kernel_size=23, sigma=(0.1, 2.0)),
])

tribyol_dataset = MinervaDataset(readers=readers, transforms=transforms)

# The resulting batch dictionary seamlessly delivers "view1", "view2", and "view3" to the model

By standardizing on dictionaries, the pipeline becomes modular and self-documenting. Complex data strategies are reduced to a simple list of operations, accelerating development and minimizing the surface area for bugs.

The dictionary-based pipeline proposed here is not a novel idea, but rather a convergence with the broader ML ecosystem. Several major frameworks have independently adopted this exact pattern. Albumentations, the de facto standard for image augmentation in computer vision, operates entirely on dictionaries, where each transform receives and returns a dict with keys such as "image", "mask", and "bboxes", routing operations to the appropriate targets by key.

MONAI, the leading framework for medical image analysis built on PyTorch, follows the same design with its dictionary-transform family, where synchronized spatial transforms across images and labels are a first-class concern. OpenMMLab's MMsegmentation pipeline passes a results dict through every stage, with each transform reading and writing by key, and a final packing step that converts the dict into the model's expected format. The fact that these projects, developed independently and targeting different domains, all converged on the same dictionary-based abstraction is strong evidence that this is the right design for flexible, composable data pipelines.

[fernandoGubiMarques]

Sounds good to me. My only further suggestion is that dataset and transform outputs should have the structure list[dict[str, Any]] instead of having a dict[str, Any]. This avoids having to deal with several dictionary keys that are related by name, but whose relationship is not enforced by the data structure. For instance, if we have the keys "view1" and "view2" in a an output dict, each subsequent transformation may need to figure out

• How many others "view%d+" there are?
• Is there a "view0"?
• Does the number even matter?
• Are there keys "label1" and "label2", etc. associated to these views?
• If there are other related keys (e.g. "info1", "info2", etc.) do they relate in the same way?
• Any further information related we may need to inform need to be manually kept track by the user (e.g. in the previous comment, we manually keep track of "view1", "view2" and "view3" to guide its "path" across transformations.).

In my proposal, each element in the list represents information that refers to a particular view. Therefore, each element in the list has the structure

{
    "data": torch.Tensor,
    "label": torch.Tensor,
    "info": some.Datatype,
}

(actually, as discussed with @GabrielBG0, only the "data" field is required, the others are optional)

Suppose we insert a "path" key to store which path across transformations this view should take. We therefore give this field the value of an integer id at some point. Suppose also that instead of a keys_to_transform init argument, each transform has a filter init argument which dictates which elements in the list it is allowed to alter. The previous example would look something like this:

transforms = _TransformsPipeline([
    # Split the original image into three separate views
    MultiViewTransform(views=3),
    AddPathField(),
    
    # Apply baseline augmentations
    RandomResizedCropTransform(size=(224, 224), filter=lambda x: x['path'] < 2),
    RandomResizedCropTransform(size=(96, 96), filter=lambda x: x['path'] == 2),
    
    # Apply asymmetric augmentations easily by targeting specific keys
    ColorJitterTransform(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1), # no filter => applied on all views.
    GaussianBlurTransform(kernel_size=23, sigma=(0.1, 2.0), filter=lambda x:x['path'] > 0),
])

Note we no longer need to keep track of each individual 'view' and 'label', etc. but only the 'path' field