New datamodules design

[djdameln]

Description

• This is a proof of concept for a new datamodule design. The base classes and all three datasets (MVTec, Folder, BTech) have been updated according to the new design.
• Fixes Support training with only normal images (no evaluation) #277

Summary of the design

• The AnomalibDataset and its subclasses now have a setup method, which is called from DataModule._setup(). When called, the dataframe will be created by using the create_dataset function. setup must be called before the dataset class can be used. This is so that we can instantiate the dataset in the constructor of the datamodule, before prepare_data is called.
• The DataModule is responsible for performing the random subset splitting. It first creates the fixed subsets in the constructor, and then performs any additional subset splitting in DataModule._setup().
• The _setup method from the DataModule is called only once, and is independent of the stage argument.
• To facilitate the dynamic subset splitting on the DataModule side, helper functions can be added to the data utils (see concatenate_datasets and random_split for an example).

Responsibilites of the different classes:

• create_dataset functions: Create a dataframe with information about the samples, including any information about the fixed train/val/test split that may follow from the folder organization or annotation files of the dataset.
• AnomalibDataset and subclasses: Prepare dataset items and ground truth labels + masks to be used as model input.
• AnomalibDataModule: Create dataloaders and perform dynamic subset splitting.

Advantages

• Clear responsibility of the different classes.
• The Dataset classes can be used standalone.
• No seeds needed, because dynamic subset splitting is done only once.

Known issues

• Some terminology is still a bit confusing, e.g. inconsistent use of val, infer and test (tranform_config_val vs pre_process_infer vs test_batch_size)
• Tests need to be updated.

[ashwinvaidya17]

I went over it and I don't have any major feedback. I feel this is a much cleaner design and more flexible.

[jpcbertoldo]

(i'm going to be a bit annoying because i think this refactor is important, and it will my work a lot)

this version is better than the other one IMO, but overall both are creating divergences with how LightningDatamodule is supposed to be used (and how their downstream Dataset is supposed to be designed)

There are some advantages in this design but I think following lightnings's patterns is a better way to go because then you make better use of it

Another guideline i'm considering: the behavior of parent classes should be kept minimal (or use the template method pattern, like lightning) to allow more composability.

i am trying to keep this short but below i give you 2 reasons to not go this way

i wrote another draft in my fork (didn't open another PR yet): my branch

• base.py
• mvtec.py (got much cleaner)

---

reason 1

the order of things is not quite right

example:

1. LightningDatamodule.prepare_data() is supposed to not change state and only manage the files in the system to make sure they`re in place
2. LightningDatamodule.setup() should be loading whatever necessary thing to the memory

In the way things are being done here (correct me if i'm wrong), the setup is being made even before the prepare_data().

well, because of (2), the Dataset somehow has to load things to the memory either (a) lazyly (following LightningDatamodule's pattern) or (b) it should be instantiated at LightningDatamodule.setup(), which means that the data preparation should not be a method (but a staticmethod could do).

in my draft i solved that by making the Dataset follow exactly the same pattern: .prepare_data() (which is calls a staticmethod) and .setup() (which will make use of the args from __init__).

---

reason 2

i think the current design is biased by the behavior in folder.py and forcing to fit the (already functional) code of the make_datate-like functions

from what i understan, a Dataset should keep the knowlege of how to find and load samples from a specific pre-defined split; it should not deal with dynamic splitting (i.e. creating random subsplits to create a validation set)

i think that kind of behavior should be at the Datamodule level because it can ensure the compatibiilty between the splits when they've all already been setup

in my draft i used torch.utils.data.Subset to do that and transfered all this behavior to AnomalibDatamodule; two advantages:

• downstream class (see mvtec.py) become cleaner
• the behavior is better factorized out

you will see that, for instance, seed and create_validation_set (smt. like that) dont need to be passed downstream to the child class (then to the function)

[jpcbertoldo]

@djdameln i'm very willing to help on this refactor in particular but unfortunately i'm struggling to find time to create a neat and well explained PR like yours

maybe can we make a call? i think it would be more productive

[djdameln]

@jpcbertoldo Thanks for your comments. As it turns out, it's quite tricky to get this design right, so it's nice to have an extra pair of eyes on this. I agree with some, but not all of your suggestions.

In the way things are being done here (correct me if i'm wrong), the setup is being made even before the prepare_data().

The call to create_dataset can be made at the end of the prepare_data method, so in this case we would create the samples after downloading the data. But yeah, I do agree that this does not really fit the intended use of prepare_data and setup.

in my draft i solved that by making the Dataset follow exactly the same pattern: .prepare_data() (which is calls a staticmethod) and .setup() (which will make use of the args from init).

I had a look at your draft and correct me if I'm wrong, but I don't think it's necessary to have the implementation of prepare_data in the dataset class. We could keep this functionality in the datamodule. That way it would be called only once for the entire datamodule, instead of three times (once for every subset).

I do see the added value of moving the setup method to the dataset class, because it allows us to instantiate the dataset object in the constructor of the datamodule. That way we can get rid of the awkward create_dataset method of my design.

you will see that, for instance, seed and create_validation_set (smt. like that) dont need to be passed downstream to the child class (then to the function)

This is true for MVTec, where the train/test set is fixed, but not for the Folder dataset where we create a random train/test split at runtime. When following your design, we would still need to ensure somehow that the same seed is used between train_data.setup and test_data.setup.

This was actually the main motivation behind my latest design. By starting with a common dataset object with the 'Full' split, we ensure that we only have to call create_samples once, so there's no need to pass the seed around.

Anyway let's continue the discussion in a call. I'll schedule one for early next week.

[jpcbertoldo]

When we perform the dynamic subset splitting in anomaly tasks, we often extract the validation set from the pre-supplied test set (for example, when val_split_mode is set to from_test). In this case, to create the validation set, we first need to create the test set and then randomly split the test set into validation and testing subsets. If we would use stage-related logic in the setup method, we would end up creating and splitting the test set in both the fit and test stages.

Hmm, ok I see. Makes sense.

This would also mean that we would have pass a seed to the subset splitting, to ensure that we'll end up with the same val/test split during validation and testing.

Got your point but still would be nice to have a seed attrribute (enforced to be not None when something random is involved, maybe?) in the datamodule, no?

Right now there is that option using the random_split but no seed being kept in the datamodule to keep a trace.
And I think it is ok to stop at the datamodule level, the datasets from the split shouldn't need to keep track of the seed IMO.

Btw, it's be nice to have an option ValSplitMode.NONE for the case where the user doesn't care about having a validation split (my case sometimes 🙃 ). (don't forget to adapt the logic in is_setup() in that case)

I feel that having to pass a seed around between different parts of the code to ensure non-overlapping subsets is error-prone and should be avoided.

Yup.

I guess the major drawback of this approach is that it leads to higher memory consumption, but since our dataset classes do not keep all of their images in memory but rather read the images from the file system when needed in __getitem__, I expect this to be negligible in most cases.

Yup, good point!

---

Are you going to make any other minor changes?

---

... decided to deviate from the intended use of the stage argument in setup....

I think it'd be nice to explicitly state this in the code and/or doc for the sake of the record and to have future devs not question themselves hahaha.

[djdameln]

Got your point but still would be nice to have a seed attrribute (enforced to be not None when something random is involved, maybe?) in the datamodule, no?

Good idea, I agree that this could be useful to have. I've added the seed argument to the datamodule. Please note that we're not using this argument when running the training from the entrypoint scripts, because we already set the global seed using seed_everything. So we don't need to pass the seed to ensure consistency between runs (but I still see the added value of the seed parameter for custom use-cases).

Btw, it's be nice to have an option ValSplitMode.NONE for the case where the user doesn't care about having a validation split (my case sometimes 🙃 ). (don't forget to adapt the logic in is_setup() in that case)

Done.

Are you going to make any other minor changes?

We're planning to merge this to a feature branch for now, as we're working on adding more functionality to the data side of the library (synthetic anomaly generation, support for video datasets), and we don't want to expose this to the main branch until it's stable. So we'll keep making incremental changes to the datamodules on the feature branch. In terms of overall design I don't expect many more changes though, so if you want to start building stuff on top of these base classes that should be fine. You could just target your PR to the feature branch.

I think it'd be nice to explicitly state this in the code and/or doc for the sake of the record and to have future devs not question themselves hahaha.

I've added an explanation to the docstring for now. We'll probably update the documentation on the datamodules at a later point.

[jpcbertoldo]

nice :)

---

So we don't need to pass the seed to ensure consistency between runs (but I still see the added value of the seed parameter for custom use-cases)

Yes, you are right.
BUT, there ways of ensuring consistency better than others haha.
Let's leave that for another chat : )

---

Thanks for the changes, great work!