Efficient way of calculating mean, std, etc of Subset Dataset

HansBambel: <p>I have a Dataset class that loads two datasets from their respective folders (train and test).<br>
I would like to create a validation set from the training set. For this I am using the <code>random_split</code> function.<br>
This results in two Subset-Datasets: <code>train_dataset</code> and <code>valid_dataset</code>. For normalization I would like to calculate the mean and std (or min/max) of the training set, but it is not possible to do a simple call of <code>mean()</code> to the dataset, because it is a Subset.<br>
What I can do is calculate the mean, std, etc on the whole <code>dataset</code>, but this would mean that the validation set is also normalized with values from the validation set. Is there an efficient way to calculate statistics of a Subset? Or do I need to loop over it?<br>
This is my code that is calculating the “wrong” mean and std (over the whole dataset)</p>
<pre><code class="lang-auto">    dataset = transportation_dataset(data_path=data_folder, train=True)
    # Split the data into training and validation set
    num_train = len(dataset)
    split_valid = int(np.floor(valid_size * num_train))
    split_train = num_train - split_valid
    train_dataset, valid_dataset = random_split(dataset, [split_train, split_valid])
    # Test dataset
    test_dataset = transportation_dataset(data_path=data_folder, train=False)

    # TODO normalize dataset (using scaler trained on training set)
    # get mean and std of trainset (for every feature)
    # Just loop over dataset?
    mean_train, std_train = torch.mean(train_dataset.dataset.data, dim=0), torch.std(train_dataset.dataset.data, dim=0)
    # mean and std is the same as (which should not be the case)
    mean_train, std_train = torch.mean(dataset.data, dim=0), torch.std(dataset.data, dim=0)
    train_dataset.dataset.data = (train_dataset.dataset.data - mean_train) / std_train
    valid_dataset.dataset.data = (valid_dataset.dataset.data - mean_train) / std_train
    test_dataset.data = (test_dataset.data - mean_train) / std_train
</code></pre>
<p>I expected to be able to do the following for calculating the mean and std for the Subsets:</p>
<pre><code class="lang-auto">    mean_train, std_train = torch.mean(train_dataset.data, dim=0), torch.std(train_dataset.data, dim=0)
    train_dataset.data = (train_datasetdata - mean_train) / std_train
    valid_dataset.data = (valid_dataset.data - mean_train) / std_train
    test_dataset.data = (test_dataset.data - mean_train) / std_train
</code></pre>

ptrblck: <p>Since you can directly access the <code>.data</code> inside your dataset, I would recommend to create the dataset indices manually and just split these indices e.g. using <code>sklearn.model_selection.train_test_split</code>.<br>
Once you have the training and validation indices, you could pass the original dataset together with the corresponding indices to a <code>Subset</code>.<br>
The statistics should then be calculated using the training subset and by indexing the <code>subset.dataset.data</code> with the training indices.</p>
<p>Let me know, if this would work for you.</p>

HansBambel: <p>Thanks for the idea! I came up with something a bit different:</p>
<pre><code class="lang-auto">    mean_train = torch.mean(train_dataset.dataset.data[train_dataset.indices], dim=0)
    std_train = torch.std(train_dataset.dataset.data[train_dataset.indices], dim=0)
    train_dataset.dataset.data[train_dataset.indices] = (train_dataset.dataset.data[train_dataset.indices] - mean_train) / std_train
    valid_dataset.dataset.data[valid_dataset.indices] = (valid_dataset.dataset.data[valid_dataset.indices] - mean_train) / std_train
    test_dataset.data = (test_dataset.data - mean_train) / std_train
</code></pre>
<p>This way I am using pytorch’s <code>random_split</code>, but calculate the stats on the train-subset.</p>

mia-minyoung: <p>Hi, thanks you for asking and sharing the answer this!</p>
<p>I encountered the same issue and tried to use your code. But I wonder if it is possible get all data with [train_dataset.indices]. I’m not sure whether it’s a list or array or a single index. If [train_dataset.indices] is a single index, two for loop are needed at the end. Is this what you do exactly?</p>

HansBambel: <p><code>train_dataset.indices</code> is a list with indices. So when I call <code>valid_dataset.dataset.data</code> I get the <strong>whole</strong> data from the dataset. In order to only get only those that belong to the Subset I can select them with indexing <code>train_dataset.dataset.data[train_dataset.indices]</code>. My data is a tensor, but it could also be a numpy array.</p>
<p>You can also loop over it I think:</p>
<pre><code class="lang-auto">current_mean = 0.0
for i in train_dataset.indices:
    current_mean += train_dataset.dataset.data[i]
current_mean /= len(train_dataset.indices)
</code></pre>
<p>Hope this helps!</p>

mia-minyoung: <p>Now I get it. Thank you !!</p>