def plot_images_with_noise(datadir, noise_values, plot_path):
"""Plot Sample MNIST images with noise."""
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))])
dataset = datasets.MNIST(datadir,
train=False,
download=True,
transform=transform)
num_noise = len(noise_values)
fig = plt.figure(figsize=(num_noise, 4))
for y in range(4):
for x in range(num_noise):
transform.transforms.append(
RandomNoise(noise_values[x], high_value=0.1307 + 2 * 0.3081))
img, _ = dataset[y]
transform.transforms.pop()
ax = fig.add_subplot(4, num_noise, y * num_noise + x + 1)
ax.set_axis_off()
ax.imshow(img.numpy().reshape((28, 28)), cmap="gray")
if y == 0:
ax.set_title("{0}%".format(noise_values[x] * 100))
plt.tight_layout()
plt.savefig(plot_path)
plt.close()
def create_test_loaders(dataset, noise_values, batch_size, data_dir):
"""
Create a list of data loaders, one for each noise value
"""
print("Creating test loaders for noise values:", noise_values)
loaders = []
for noise in noise_values:
transform_noise_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
(0.50707516, 0.48654887, 0.44091784),
(0.26733429, 0.25643846, 0.27615047),
),
RandomNoise(noise,
high_value=0.5 + 2 * 0.20,
low_value=0.5 - 2 * 0.2),
])
testset = getattr(datasets, dataset)(root=data_dir,
train=False,
transform=transform_noise_test)
loaders.append(
DataLoader(testset, batch_size=batch_size, shuffle=False))
return loaders
def run_noise_tests(self):
"""
Test the model with different noise values and return test metrics.
"""
ret = {}
# Test with noise
for noise in [
0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5
]:
self.logger.info("Noise: %s", noise)
transform = transforms.Compose([
transforms.ToTensor(),
RandomNoise(noise, high_value=0.1307 + 2 * 0.3081),
transforms.Normalize((0.1307, ), (0.3081, )),
])
test_loader = torch.utils.data.DataLoader(
datasets.MNIST(self.data_dir, train=False,
transform=transform),
batch_size=self.test_batch_size,
shuffle=True,
)
ret[noise] = self.test(test_loader)
return ret
def get_test_dataset(self, noise_level=0.0):
if noise_level not in self.test_datasets:
all_transforms = [transforms.ToTensor()]
if noise_level > 0.0:
all_transforms.append(RandomNoise(noise_level))
all_transforms.append(transforms.Normalize((MEAN, ), (STDEV, )))
transform = transforms.Compose(all_transforms)
self.test_datasets[noise_level] = CachingDataset(
datasets.MNIST(self.folder,
train=False,
download=True,
transform=transform))
return self.test_datasets[noise_level]
def set_noise_loader(self, noise):
"""Defines noise loader"""
noise_transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize(self.stats_mean, self.stats_std),
RandomNoise(noise, high_value=0.5 + 2 * 0.20, low_value=0.5 - 2 * 0.2),
]
)
noise_set = self.dataset(
root=self.data_dir, train=False, transform=noise_transform
)
self.noise_loader = DataLoader(
dataset=noise_set, batch_size=self.batch_size_test, shuffle=False
)
def _configure_dataloaders(self, config):
data_dir = config["data_dir"]
batch_size = config["batch_size"]
test_batch_size = config["test_batch_size"]
first_epoch_batch_size = config["first_epoch_batch_size"]
transform = [
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
# transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
]
# Add noise
noise = config.get("noise", 0.0)
if noise > 0.0:
transform.append(
RandomNoise(noise,
high_value=0.5 + 2 * 0.2,
low_value=0.5 - 2 * 0.2))
train_dataset = datasets.CIFAR10(
data_dir, train=True, transform=transforms.Compose(transform))
test_dataset = datasets.CIFAR10(data_dir,
train=False,
transform=transforms.Compose(
transform[2:]))
self.train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True)
self.test_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=test_batch_size, shuffle=False)
if first_epoch_batch_size != batch_size:
self.first_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=first_epoch_batch_size, shuffle=True)
else:
self.first_loader = self.train_loader
def create_test_loaders(noise_values, batch_size, data_dir):
"""Create a list of data loaders, one for each noise value."""
print("Creating test loaders for noise values:", noise_values)
loaders = []
for noise in noise_values:
transform_noise_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
RandomNoise(noise,
high_value=0.5 + 2 * 0.20,
low_value=0.5 - 2 * 0.2),
])
testset = datasets.CIFAR10(root=data_dir,
train=False,
transform=transform_noise_test)
loaders.append(
DataLoader(testset, batch_size=batch_size, shuffle=False))
return loaders
def __init__(self, config=None):
defaults = dict(
dataset_name=None,
data_dir=None,
batch_size_train=128,
batch_size_test=128,
stats_mean=None,
stats_std=None,
augment_images=False,
test_noise=False,
noise_level=0.1,
)
defaults.update(config)
self.__dict__.update(defaults)
# expand ~
self.data_dir = os.path.expanduser(self.data_dir)
# recover mean and std to normalize dataset
if not self.stats_mean or not self.stats_std:
tempset = getattr(datasets, self.dataset_name)(
root=self.data_dir, train=True, transform=transforms.ToTensor()
)
self.stats_mean = (tempset.data.float().mean().item() / 255,)
self.stats_std = (tempset.data.float().std().item() / 255,)
del tempset
# set up transformations
transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize(self.stats_mean, self.stats_std),
]
)
if not self.augment_images:
aug_transform = transform
else:
aug_transform = transforms.Compose(
[
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(self.stats_mean, self.stats_std),
]
)
# load train set
train_set = getattr(datasets, self.dataset_name)(
root=self.data_dir, train=True, transform=aug_transform
)
self.train_loader = DataLoader(
dataset=train_set, batch_size=self.batch_size_train, shuffle=True
)
# load test set
test_set = getattr(datasets, self.dataset_name)(
root=self.data_dir, train=False, transform=transform
)
self.test_loader = DataLoader(
dataset=test_set, batch_size=self.batch_size_test, shuffle=False
)
# noise dataset
if self.test_noise:
noise = self.noise_level
noise_transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize(self.stats_mean, self.stats_std),
RandomNoise(
noise, high_value=0.5 + 2 * 0.20, low_value=0.5 - 2 * 0.2
),
]
)
noise_set = getattr(datasets, self.dataset_name)(
root=self.data_dir, train=False, transform=noise_transform
)
self.noise_loader = DataLoader(
dataset=noise_set, batch_size=self.batch_size_test, shuffle=False
)
def load_from_torch_vision(self):
# special dataloader case
if isinstance(self.batch_size_train, Iterable) or isinstance(
self.batch_size_test, Iterable):
dataloader_type = VaryingDataLoader
else:
dataloader_type = DataLoader
# expand ~
self.data_dir = os.path.expanduser(self.data_dir)
# recover mean and std to normalize dataset
if not self.stats_mean or not self.stats_std:
tempset = getattr(datasets, self.dataset_name)(
root=self.data_dir,
train=True,
transform=transforms.ToTensor())
self.stats_mean = (tempset.data.float().mean().item() / 255, )
self.stats_std = (tempset.data.float().std().item() / 255, )
del tempset
# set up transformations
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(self.stats_mean, self.stats_std),
])
# set up augment transforms for training
if not self.augment_images:
aug_transform = transform
else:
aug_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(self.stats_mean, self.stats_std),
])
# load train set
train_set = getattr(datasets,
self.dataset_name)(root=self.data_dir,
train=True,
transform=aug_transform)
self.train_loader = dataloader_type(dataset=train_set,
batch_size=self.batch_size_train,
shuffle=True)
# load test set
test_set = getattr(datasets, self.dataset_name)(root=self.data_dir,
train=False,
transform=transform)
self.test_loader = dataloader_type(dataset=test_set,
batch_size=self.batch_size_test,
shuffle=False)
# noise dataset
if self.test_noise:
noise = self.noise_level
noise_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(self.stats_mean, self.stats_std),
RandomNoise(noise,
high_value=0.5 + 2 * 0.20,
low_value=0.5 - 2 * 0.2),
])
noise_set = getattr(datasets,
self.dataset_name)(root=self.data_dir,
train=False,
transform=noise_transform)
self.noise_loader = dataloader_type(
dataset=noise_set,
batch_size=self.batch_size_test,
shuffle=False)