How to calculate mean and std for 1 channel using datalodaer

nima_pw: <p>Hi,</p>
<p>I have a code that calculates <code>mean and std</code> of my image dataset through dataloader:</p>
<pre><code class="lang-auto">tfms = transforms.Compose([
    transforms.Resize((128, 128)),
                  transforms.ToTensor(),
                  transforms.Normalize(mean = (0, 0, 0), std  = (1, 1, 1))])

train_files = glob.glob(train_dir + '/*.jpg')
train_ds = MyDataset(train_files , transform = tfms)
train_dl = torch.utils.data.DataLoader(train_ds, 
                                       batch_size=16, 
                                       shuffle=True,)

####### COMPUTE MEAN / STD

# placeholders
psum    = torch.tensor([0.0, 0.0, 0.0]).to(device)
psum_sq = torch.tensor([0.0, 0.0, 0.0]).to(device)

# loop through images
for inputs, _ in tqdm(train_dl):
    inputs = inputs.to(device)
    psum    += inputs.sum(axis        = [0, 2, 3])
    psum_sq += (inputs ** 2).sum(axis = [0, 2, 3])


# ####### FINAL CALCULATIONS

# pixel count
print(len(train_ds))
count = len(train_ds) * resize * resize

# mean and std
total_mean = psum / count
total_var  = (psum_sq / count) - (total_mean ** 2)
total_std  = torch.sqrt(total_var)

# output
print('mean: '  + str(total_mean))
print('std:  '  + str(total_std))
</code></pre>
<p>I want to add <code>transforms.Lambda(noise_extract)</code> to <code>transforms.Compose</code> that extract noise residual and converts images from 3 channels to 1 channel:</p>
<pre><code class="lang-auto">tfms = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.Lambda(noise_extract),
                  transforms.ToTensor(),
                  transforms.Normalize(mean = (0, 0, 0), std  = (1, 1, 1))])
</code></pre>
<p>I want to know how I must change the above code for calculating <code>mean &amp; std</code> for the new <code>one channel</code> output of the second <code>tfms</code>?</p>

nima_pw: <p>How do I change the above code to compute <code>mean &amp; std</code> for one channel grayscale?</p>

thecho7: <p>Since a grayscale image has only 1 channel, modify the normalize function like below</p>
<pre><code class="lang-auto">tfms = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.Lambda(noise_extract),
                  transforms.ToTensor(),
                  transforms.Normalize(mean = (0,), std  = (1,))])
</code></pre>
<p>And <code>psum</code>, <code>psum_sq</code> should be modified.</p>
<pre><code class="lang-auto">psum    = torch.tensor([0.0]).to(device)
psum_sq = torch.tensor([0.0]).to(device)
</code></pre>

nima_pw: <aside class="quote no-group" data-username="nima_pw" data-post="1" data-topic="175656">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/n/e19adc/48.png" class="avatar"> nima_pw:</div>
<blockquote>
<p><code>[0, 2, 3]</code></p>
</blockquote>
</aside>
<p>Dear <a class="mention" href="/u/thecho7">@thecho7</a>, thanks for your reply.<br>
How about <code>[0, 2, 3]</code>? I don’t know what it is, and should I change it?</p>
<pre><code>psum    += inputs.sum(axis        = [0, 2, 3])
psum_sq += (inputs ** 2).sum(axis = [0, 2, 3])
</code></pre>

thecho7: <p><code>axis</code> or <code>dim</code> means to sum up along the given dimensions.</p>
<pre><code class="lang-auto">a = torch.rand(2, 3, 4, 4)
b = a.sum(dim=[0, 2, 3])
</code></pre>
<p>It means b is a 1-dimension with size 3 tensor.</p>
<pre><code class="lang-auto">b[0] = a[0, 0, 0, 0] + a[0, 0, 0, 1] + ... + a[0, 0, 3, 3]
     + a[1, 0, 0, 0] + a[1, 0, 0, 1] + ... + a[1, 0, 3, 3]

b[1] = a[0, 1, 0, 0] + a[0, 1, 0, 1] + ... + a[0, 1, 3, 3]
     + a[1, 1, 0, 0] + a[1, 1, 0, 1] + ... + a[1, 1, 3, 3]

b[2] = a[0, 2, 0, 0] + a[0, 2, 0, 1] + ... + a[0, 2, 3, 3]
     + a[1, 2, 0, 0] + a[1, 2, 0, 1] + ... + a[1, 2, 3, 3]
</code></pre>