Save MNIST dataset with added noise-CodePudding

I want to save the new MNIST dataset tensors after adding noise.

mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())

def add_noise(dataset):
  for data in dataset:
      img, _ = data[0], data[1]
      noisy_data = torch.tensor(random_noise(img, mode='gaussian', mean=0, var=0.05, clip=True))
      return noisy_data

train_gauss = add_noise(mnist_trainset)

The code above only saves one image of the MNIST dataset, how can I save all tensors at once?

CodePudding user response：

Your return is inside the for loop, so at the end of the first iteration it already returns the noisy_data.

You should take it out of the for loop and then you could use a list to return the whole dataset at once, like so:

mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())

def add_noise(dataset):
  noisy_data = []
  for data in dataset:
      img, _ = data[0], data[1]
      noisy_data  = torch.tensor(random_noise(img, mode='gaussian', mean=0, var=0.05, clip=True))
  return noisy_data

train_gauss = add_noise(mnist_trainset)