normalize_CIFAR10.py

import matplotlib.pyplot as plt
import numpy as np
import os
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.prune as prune
import torch.optim as optim
import torchvision
import torchvision.models as models
import torchvision.transforms as transforms
from torch import nn
from torch.nn import Conv2d,BatchNorm2d,ReLU,MaxPool2d
#from tqdm import tqdm
from tqdm import tqdm_notebook as tqdm
%matplotlib inline


def get_data_set(train_flag=True):
    if train_flag:
        data_set = torchvision.datasets.__dict__['CIFAR10'](root='./dataset', train=True, download=True)
    else:
        data_set = torchvision.datasets.__dict__['CIFAR10'](root='./dataset', train=False,download=True)    
    return data_set
  
train_data_set = get_data_set(train_flag=True)
test_data_set = get_data_set(train_flag=False)

print(train_data_set.data.mean(axis=(0,1,2))/255)
print(train_data_set.data.std(axis=(0,1,2))/255)

print(test_data_set.data.mean(axis=(0,1,2))/255)
print(test_data_set.data.std(axis=(0,1,2))/255)