<p>The solution I found is to create three data sets since the beginning with the correct transforms you want. Then you have 3 data set objects and you give it to the <code>torch.utils.data.Subset(train_dataset, train_indices)</code>. The crux is this essentially:</p>
<pre><code>    # load the dataset
    path_to_data_set: str = str(Path(path_to_data_set).expanduser())
    train_dataset = datasets.MNIST(root=path_to_data_set, train=True,
                                   download=True, transform=train_transform)
    val_dataset = datasets.MNIST(root=path_to_data_set, train=True,
                                 download=True, transform=val_transform)
    indices = list(range(len(train_dataset)))
    train_indices, val_indices = split_inidices(indices, test_size=val_size, random_state=seed, shuffle=shuffle)
    train_dataset = torch.utils.data.Subset(train_dataset, train_indices)
    val_dataset = torch.utils.data.Subset(val_dataset, val_indices)
    train_loader, val_loader = get_serial_or_distributed_dataloaders(train_dataset,
                                                                     val_dataset,
                                                                     batch_size,
                                                                     batch_size_eval,
                                                                     rank,
                                                                     world_size,
                                                                     merge,
                                                                     num_workers,
                                                                     pin_memory
                                                                     )
</code></pre>
<p>then you can create whatever dataloaders you want later. This way you don't have to change the transform in the first place.</p>
<hr />
<p>full code:</p>
<pre><code>&quot;&quot;&quot;

# - data augmentation
Current belief is that augmenting the validation set should be fine, especially if you want to actually encourage
generalization since it makes the val set harder and it allows you to make val split percentage slightly lower since
your validation set was increased size.

For reproducibility of other work, especially for scientific pursues rather than &quot;lets beat state of the art&quot; - to make
it easier to compare results use what they use. e.g. it seems only augmenting the train set is the common thing,
especially when I looked at the augmentation strategies in min-imagenet and mnist.

Test set augmentation helps mostly to make test set harder (so acc should go down) - but it also increases variance
since the data size was increased. If you are reporting results most likely augmenting the data set is a good idea
- especially if you are going to compute test set errors when comparing accuracy with previous work.
Also, the way CI intervals are computed with t_p * std_n / sqrt n, means that the avg test error will be smaller, so
you are fine in general.
Default code I see doesn't augment test set so I most likely won't either.

ref:
    - https://stats.stackexchange.com/questions/320800/data-augmentation-on-training-set-only/320967#320967
    - https://arxiv.org/abs/1809.01442, https://stats.stackexchange.com/a/390470/28986

# - pin_memory
For data loading, passing pin_memory=True to a DataLoader will automatically put the fetched data Tensors in pinned
memory, and thus enables faster data transfer to CUDA-enabled GPUs. Note on pinning:
This is an advanced tip. If you overuse pinned memory, it can cause serious problems when running low on RAM, and
you should be aware that pinning is often an expensive operation. Thus, will leave it's default as False.

ref:
    - on pin_memory: https://pytorch.org/docs/stable/data.html
&quot;&quot;&quot;
from typing import Callable, Optional, Union

import numpy as np
import torch
from numpy.random import RandomState
from torch.utils.data import Dataset, SubsetRandomSampler, random_split, DataLoader, RandomSampler


def get_train_val_split_random_sampler(
        train_dataset: Dataset,
        val_dataset: Dataset,
        val_size: float = 0.2,
        batch_size: int = 128,
        batch_size_eval: int = 64,
        num_workers: int = 4,
        pin_memory: bool = False
        #     random_seed: Optional[int] = None,
) -&gt; tuple[DataLoader, DataLoader]:
    &quot;&quot;&quot;
    Note:
        - this will use different transforms for val and train if the objects you pass have different transforms.
        - note train_dataset, val_dataset whill often be the same data set object but different instances with different
        transforms for each data set.

    Recommended use:
        - this one is recommended when you want the train &amp; val to have different transforms e.g. when doing scientific
        work - instead of beating benchmark work - and the train, val sets had different transforms.

    ref:
        - https://gist.github.com/MattKleinsmith/5226a94bad5dd12ed0b871aed98cb123
    &quot;&quot;&quot;
    assert 0 &lt;= val_size &lt;= 1.0, f&quot;Error: {val_size} valid_size should be in the range [0, 1].&quot;
    num_train = len(train_dataset)
    indices = list(range(num_train))
    split_idx = int(np.floor(val_size * num_train))

    # I don't think this is needed since later the sampler randomly samples data from a given list
    # if shuffle == True:
    #     np.random.seed(random_seed)
    #     np.random.shuffle(indices)

    train_idx, valid_idx = indices[:split_idx], indices[split_idx:]
    assert len(train_idx) != 0 and len(valid_idx) != 0

    # Samples elements randomly from a given list of indices, without replacement.
    train_sampler = SubsetRandomSampler(train_idx)
    valid_sampler = SubsetRandomSampler(valid_idx)

    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batch_size, sampler=train_sampler,
                                               num_workers=num_workers, pin_memory=pin_memory)
    valid_loader = torch.utils.data.DataLoader(val_dataset,
                                               batch_size=batch_size_eval, sampler=valid_sampler,
                                               num_workers=num_workers, pin_memory=pin_memory)
    return train_loader, valid_loader


def get_train_val_split_with_split(
        train_dataset: Dataset,
        train_val_split: list[int, int],  # e.g. [50_000, 10_000] for mnist
        batch_size: int = 128,
        batch_size_eval: int = 64,
        num_workers: int = 4,
        pin_memory: bool = False
) -&gt; tuple[DataLoader, DataLoader]:
    &quot;&quot;&quot;
    Note:
        - this will have the train and val sets have the same transform.

    ref:
        - https://gist.github.com/MattKleinsmith/5226a94bad5dd12ed0b871aed98cb123
        - change transform: https://discuss.pytorch.org/t/changing-transforms-after-creating-a-dataset/64929/4
    &quot;&quot;&quot;
    train_dataset, valid_dataset = random_split(train_dataset, train_val_split)
    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batch_size, num_workers=num_workers, pin_memory=pin_memory)
    valid_loader = torch.utils.data.DataLoader(valid_dataset,
                                               batch_size=batch_size_eval, num_workers=num_workers,
                                               pin_memory=pin_memory)
    return train_loader, valid_loader


def get_serial_or_distributed_dataloaders(train_dataset: Dataset,
                                          val_dataset: Dataset,
                                          batch_size: int = 128,
                                          batch_size_eval: int = 64,
                                          rank: int = -1,
                                          world_size: int = 1,
                                          merge: Optional[Callable] = None,
                                          num_workers: int = -1,  # -1 means its running serially
                                          pin_memory: bool = False,
                                          ):
    &quot;&quot;&quot;

    &quot;&quot;&quot;
    from uutils.torch_uu.distributed import is_running_serially
    if is_running_serially(rank):
        train_sampler = RandomSampler(train_dataset)
        val_sampler = RandomSampler(val_dataset)
        num_workers = 4 if num_workers == -1 else num_workers
    else:
        assert (batch_size &gt;= world_size), f'Each worker must get at least one data point, so batch_size &gt;= world_size' \
                                           f'but got: {batch_size}{world_size}'
        from torch.utils.data import DistributedSampler

        # note: shuffle = True by default
        train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
        val_sampler = DistributedSampler(val_dataset, num_replicas=world_size, rank=rank)
        # set the input num_workers but for ddp 0 is recommended afaik, todo - check
        num_workers = 0 if num_workers == -1 else num_workers
    # get dist dataloaders
    train_loader = DataLoader(train_dataset,
                              batch_size=batch_size,
                              sampler=train_sampler,
                              collate_fn=merge,
                              num_workers=num_workers,
                              pin_memory=pin_memory)
    val_loader = DataLoader(val_dataset,
                            batch_size=batch_size_eval,
                            sampler=val_sampler,
                            collate_fn=merge,
                            num_workers=num_workers,
                            pin_memory=pin_memory)
    # return dataloaders
    # dataloaders = {'train': train_dataloader, 'val': val_dataloader, 'test': test_dataloader}
    # iter(train_dataloader)  # if this fails its likely your running in pycharm and need to set num_workers flag to 0
    return train_loader, val_loader


def split_inidices(indices: list,
                   test_size: Optional = None,
                   random_state: Optional[Union[int, RandomState, None]] = None,
                   shuffle: bool = False,  # false for reproducibility, and any split is as good as any other.
                   ) -&gt; tuple[list[int], list[int]]:
    import sklearn.model_selection
    # - api: https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html
    train_indices, val_indices = sklearn.model_selection.train_test_split(indices, test_size=test_size,
                                                                          random_state=random_state,
                                                                          shuffle=shuffle)
    return train_indices, val_indices

# - visualization help


&quot;&quot;&quot;
Inspired from:
    - https://gist.github.com/MattKleinsmith/5226a94bad5dd12ed0b871aed98cb123
    - https://www.geeksforgeeks.org/training-neural-networks-with-validation-using-pytorch/
&quot;&quot;&quot;
from argparse import Namespace
from pathlib import Path
from typing import Optional, Callable

import numpy as np
import torch
from torch.utils.data import random_split, DataLoader
from torchvision import datasets
from torchvision.transforms import transforms

from uutils.torch_uu.dataloaders.common import split_inidices, \
    get_serial_or_distributed_dataloaders

NORMALIZE_MNIST = transforms.Normalize((0.1307,), (0.3081,))  # MNIST


def get_train_valid_test_data_loader_helper_for_mnist(args: Namespace) -&gt; dict:
    train_kwargs = {'path_to_data_set': args.path_to_data_set,
                    'batch_size': args.batch_size,
                    'batch_size_eval': args.batch_size_eval,
                    'augment_train': args.augment_train,
                    'augment_val': args.augment_val,
                    'num_workers': args.num_workers,
                    'pin_memory': args.pin_memory,
                    'rank': args.rank,
                    'world_size': args.world_size,
                    'merge': None
                    }
    test_kwargs = {'path_to_data_set': args.path_to_data_set,
                   'batch_size_eval': args.batch_size_eval,
                   'augment_test': args.augment_train,
                   'num_workers': args.num_workers,
                   'pin_memory': args.pin_memory,
                   'rank': args.rank,
                   'world_size': args.world_size,
                   'merge': None
                   }
    train_loader, val_loader = get_train_valid_loader(**train_kwargs)
    test_loader: DataLoader = get_test_loader(**test_kwargs)
    dataloaders: dict = {'train': train_loader, 'val': val_loader, 'test': test_loader}
    return dataloaders


def get_transform(augment: bool):
    if augment:
        transform = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            NORMALIZE_MNIST
        ])
    else:
        transform = transforms.Compose([
            transforms.ToTensor(),
            NORMALIZE_MNIST
        ])
    return transform


def get_train_valid_loader(path_to_data_set: Path,
                           batch_size: int = 128,
                           batch_size_eval: int = 64,
                           seed: Optional[int] = None,
                           augment_train: bool = True,
                           augment_val: bool = False,
                           val_size: Optional[float] = 0.2,
                           shuffle: bool = False,  # false for reproducibility, and any split is as good as any other.
                           num_workers: int = -1,
                           pin_memory: bool = False,

                           rank: int = -1,
                           world_size: int = 1,
                           merge: Optional[Callable] = None,
                           ) -&gt; tuple[DataLoader, DataLoader]:
    &quot;&quot;&quot;
    Utility function for loading and returning train and valid
    multi-process iterators over the MNIST dataset. A sample
    9x9 grid of the images can be optionally displayed.
    If using CUDA, num_workers should be set to 1 and pin_memory to True.
    &quot;&quot;&quot;
    # train_kwargs = {'batch_size': args.batch_size}

    # define transforms
    train_transform = get_transform(augment_train)
    val_transform = get_transform(augment_val)

    # load the dataset
    path_to_data_set: str = str(Path(path_to_data_set).expanduser())
    train_dataset = datasets.MNIST(root=path_to_data_set, train=True,
                                   download=True, transform=train_transform)
    val_dataset = datasets.MNIST(root=path_to_data_set, train=True,
                                 download=True, transform=val_transform)
    indices = list(range(len(train_dataset)))
    train_indices, val_indices = split_inidices(indices, test_size=val_size, random_state=seed, shuffle=shuffle)
    train_dataset = torch.utils.data.Subset(train_dataset, train_indices)
    val_dataset = torch.utils.data.Subset(val_dataset, val_indices)
    train_loader, val_loader = get_serial_or_distributed_dataloaders(train_dataset,
                                                                     val_dataset,
                                                                     batch_size,
                                                                     batch_size_eval,
                                                                     rank,
                                                                     world_size,
                                                                     merge,
                                                                     num_workers,
                                                                     pin_memory
                                                                     )
    return train_loader, val_loader


def get_test_loader(path_to_data_set,
                    batch_size_eval: int = 64,
                    shuffle: bool = True,
                    augment_test: bool = False,
                    num_workers: int = -1,
                    pin_memory=False,

                    rank: int = -1,
                    world_size: int = 1,
                    merge: Optional[Callable] = None,
                    ) -&gt; DataLoader:
    &quot;&quot;&quot;
    Utility function for loading and returning a multi-process
    test iterator over the MNIST dataset.
    If using CUDA, num_workers should be set to 1 and pin_memory to True.
    Params
    ------
    - path_to_data_set: path directory to the dataset.
    - batch_size: how many samples per batch to load.
    - shuffle: whether to shuffle the dataset after every epoch.
    - num_workers: number of subprocesses to use when loading the dataset.
    - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to
      True if using GPU.
    Returns
    -------
    - data_loader: test set iterator.

    Note:
        - it knows it's the test set since train=False in the body when creating the data set.
    &quot;&quot;&quot;
    # define transform
    test_transform = get_transform(augment_test)

    # load the dataset
    path_to_data_set: str = str(Path(path_to_data_set).expanduser())
    test_dataset = datasets.MNIST(root=path_to_data_set,
                                  train=False,  # ensures its test set
                                  download=True,
                                  transform=test_transform)
    _, test_loader = get_serial_or_distributed_dataloaders(test_dataset,
                                                            test_dataset,
                                                            batch_size_eval,
                                                            batch_size_eval,
                                                            rank,
                                                            world_size,
                                                            merge,
                                                            num_workers,
                                                            pin_memory,
                                                            )
    return test_loader
</code></pre>
<p>repo that came from this with permanent github link: <a href="https://github.com/brando90/ultimate-utils/blob/ef2217c07b43aa5354f7b6f8f1761c5f16017874/ultimate-utils-proj-src/uutils/torch_uu/dataloaders/mnist.py#L22" rel="nofollow noreferrer">https://github.com/brando90/ultimate-utils/blob/ef2217c07b43aa5354f7b6f8f1761c5f16017874/ultimate-utils-proj-src/uutils/torch_uu/dataloaders/mnist.py#L22</a></p>
<hr />
<p>related:</p>
<ul>
<li><a href="https://discuss.pytorch.org/t/changing-transformation-applied-to-data-during-training/15671/14" rel="nofollow noreferrer">https://discuss.pytorch.org/t/changing-transformation-applied-to-data-during-training/15671/14</a></li>
<li><a href="https://discuss.pytorch.org/t/changing-transforms-after-creating-a-dataset/64929/7" rel="nofollow noreferrer">https://discuss.pytorch.org/t/changing-transforms-after-creating-a-dataset/64929/7</a></li>
<li><a href="https://discuss.pytorch.org/t/apply-different-transform-data-augmentation-to-train-and-validation/63580/13" rel="nofollow noreferrer">https://discuss.pytorch.org/t/apply-different-transform-data-augmentation-to-train-and-validation/63580/13</a></li>
</ul>
What is the proper way to create training, validation and test set in pytorch or change the transform of an already created data set?
<p>I noticed that a standard thing like getting the validation set in PyTorch is not as common as one would expect or not obviously available in the pytorch library.</p>
<p>I found two websites that do it their own way:
- <a href="https://gist.github.com/MattKleinsmith/5226a94bad5dd12ed0b871aed98cb123" rel="nofollow noreferrer">https://gist.github.com/MattKleinsmith/5226a94bad5dd12ed0b871aed98cb123</a>
- <a href="https://www.geeksforgeeks.org/training-neural-networks-with-validation-using-pytorch/" rel="nofollow noreferrer">https://www.geeksforgeeks.org/training-neural-networks-with-validation-using-pytorch/</a></p>
<p>but they have their problems because the second one force you to have both the train &amp; validation set have the same transforms and the first one splits with respect to the data loader - which is then impossible to give easily to a distributed data loader afaik.</p>
<p>If that is not the way to do it then what is the right proper way to create a train, val and test set?</p>