r_loss = vae_criterion(
outputs, inputs, size_average=False) / inputs.size(0)
except RuntimeError as e:
logger.error('Runtime error. This could possibly be due to using '
'the wrong encoder / decoder for this dataset. '
'If you are using MNIST, for example, use the '
'arguments `--encoder_type mnist --decoder_type '
'mnist`')
raise e
kl = (0.5 * (vae.std**2 + vae.mu**2 - 2. * torch.log(vae.std) -
1.).sum(1).mean())
msssim = ms_ssim(inputs, outputs)
self.losses.vae = (r_loss + beta_kld * kl)
self.results.update(KL_divergence=kl.item(), ms_ssim=msssim.item())
def visualize(self, inputs, targets, Z):
vae = self.nets.vae
outputs = vae(inputs)
self.add_image(outputs, name='reconstruction')
self.add_image(inputs, name='ground truth')
self.add_scatter(vae.mu.data, labels=targets.data, name='latent values')
self.decoder.visualize(Z)
register_plugin(VAE)
"""
for _ in range(discriminator_updates):
self.data.next()
inputs, Z = self.inputs('inputs', 'Z')
generated = self.generator.generate(Z)
self.discriminator.routine(inputs, generated.detach())
self.optimizer_step()
self.penalty.train_step()
for _ in range(generator_updates):
self.generator.train_step()
def eval_step(self):
self.data.next()
inputs, Z = self.inputs('inputs', 'Z')
generated = self.generator.generate(Z)
self.discriminator.routine(inputs, generated)
self.penalty.routine(auto_input=True)
self.generator.routine(auto_input=True)
def visualize(self, images, Z):
self.add_image(images, name='ground truth')
generated = self.generator.generate(Z)
self.discriminator.visualize(images, generated)
self.generator.visualize(Z)
register_plugin(GAN)
inferred = self.bidirectional_model.encoder.encode(inputs)
self.discriminator.routine(
inputs, generated.detach(), inferred.detach(), Z)
self.optimizer_step()
self.penalty.routine((inputs, inferred))
self.optimizer_step()
self.bidirectional_model.train_step()
def eval_step(self):
self.data.next()
inputs, Z = self.inputs('inputs', 'Z')
generated = self.bidirectional_model.decoder.decode(Z)
inferred = self.bidirectional_model.encoder.encode(inputs)
self.discriminator.routine(inputs, generated, inferred, Z)
self.bidirectional_model.eval_step()
def visualize(self, inputs, Z, targets):
generated = self.bidirectional_model.decoder.decode(Z)
inferred = self.bidirectional_model.encoder.encode(inputs)
self.bidirectional_model.visualize()
self.discriminator.visualize(inputs, generated, inferred, Z, targets)
register_plugin(ALI)
classifier_type='convnet',
classifier_args=dict(dropout=0.2),
Encoder=None):
'''Builds a simple image classifier.
Args:
classifier_type (str): Network type for the classifier.
classifier_args: Classifier arguments. Can include dropout,
batch_norm, layer_norm, etc.
'''
classifier_args = classifier_args or {}
shape = self.get_dims('x', 'y', 'c')
dim_a = self.get_dims('attributes')
Encoder_, args = update_encoder_args(shape,
model_type=classifier_type,
encoder_args=classifier_args)
Encoder = Encoder or Encoder_
args.update(**classifier_args)
classifier = Encoder(shape, dim_out=dim_a, **args)
self.nets.classifier = classifier
register_plugin(ImageClassification)
register_plugin(ImageAttributeClassification)
def build(self, dim_z=64):
self.encoder.build(dim_out=dim_z)
self.decoder.build(dim_in=dim_z)
encoder = self.nets.encoder
decoder = self.nets.decoder
ae = AENetwork(encoder, decoder)
self.nets.ae = ae
def routine(self, inputs, targets, ae_criterion=F.mse_loss):
'''
Args:
ae_criterion: Criterion for the autoencoder.
'''
ae = self.nets.ae
outputs = ae(inputs)
r_loss = ae_criterion(
outputs, inputs, size_average=False) / inputs.size(0)
self.losses.ae = r_loss
def visualize(self, inputs, targets):
ae = self.nets.ae
outputs = ae(inputs)
self.add_image(outputs, name='reconstruction')
self.add_image(inputs, name='ground truth')
register_plugin(Autoencoder)
"""
for _ in range(discriminator_updates):
self.data.next()
inputs, Z = self.inputs('inputs', 'Z')
generated = self.generator.generate(Z)
self.discriminator.routine(inputs, generated.detach())
self.optimizer_step()
self.penalty.train_step()
for _ in range(generator_updates):
self.generator.train_step()
def eval_step(self):
self.data.next()
inputs, Z = self.inputs('inputs', 'Z')
generated = self.generator.generate(Z)
self.discriminator.routine(inputs, generated)
self.penalty.routine(auto_input=True)
self.generator.routine(auto_input=True)
def visualize(self, images, Z):
self.add_image(images, name='ground truth')
generated = self.generator.generate(Z)
self.discriminator.visualize(images, generated)
self.generator.visualize(Z)
register_plugin(GAN)
try:
uniques.remove(-1)
except ValueError:
pass
dim_l = len(uniques)
else:
if len(train_set.train_data.shape) == 4:
dim_c, dim_x, dim_y = train_set[0][0].size()
else:
dim_x, dim_y = train_set[0][0].size()
dim_c = 1
labels = train_set.train_labels
if not isinstance(labels, list):
labels = labels.numpy()
dim_l = len(np.unique(labels))
dims = dict(x=dim_x, y=dim_y, c=dim_c, labels=dim_l)
input_names = ['images', 'targets', 'index']
self.add_dataset('train', train_set)
self.add_dataset('test', test_set)
self.set_input_names(input_names)
self.set_dims(**dims)
if scale is not None:
self.set_scale(scale)
register_plugin(TorchvisionDatasetPlugin)
train=dict(epochs=200, save_on_best='losses.classifier'))
def build(self, classifier_type='convnet',
classifier_args=dict(dropout=0.2), Encoder=None):
'''Builds a simple image classifier.
Args:
classifier_type (str): Network type for the classifier.
classifier_args: Classifier arguments. Can include dropout,
batch_norm, layer_norm, etc.
'''
classifier_args = classifier_args or {}
shape = self.get_dims('x', 'y', 'c')
dim_a = self.get_dims('attributes')
Encoder_, args = update_encoder_args(
shape, model_type=classifier_type, encoder_args=classifier_args)
Encoder = Encoder or Encoder_
args.update(**classifier_args)
classifier = Encoder(shape, dim_out=dim_a, **args)
self.nets.classifier = classifier
register_plugin(ImageClassification)
register_plugin(ImageAttributeClassification)
transform=transform,
download=True)
test_set = Dataset(data_path,
train=False,
transform=test_transform,
download=True)
return train_set, test_set
def _handle_ImageFolder(self,
Dataset,
data_path,
transform=None,
test_transform=None,
**kwargs):
train_set = Dataset(f'{data_path}/train', transform=transform)
if 'UNL' in train_set.classes:
train_set.targets = np.array(train_set.targets)
train_set.targets[train_set.targets ==
train_set.class_to_idx['UNL']] = -1
train_set.class_to_idx['UNL'] = -1
val_set = Dataset(f'{data_path}/val', transform=test_transform)
test_set = Dataset(f'{data_path}/test', transform=test_transform)
return train_set, val_set, test_set
# Removing all registered DatasetPlugins
_PLUGINS.clear()
register_plugin(SSLDatasetPlugin)
DataHandler.make_iterator = make_iterator
DataHandler.__next__ = __next__
self.add_dataset('test', test_set)
self.set_input_names(input_names)
self.set_dims(**dims)
self.set_scale((-1, 1))
def make_split(self, data_path, split, Dataset, train_transform,
test_transform):
train_set = Dataset(root=data_path, transform=train_transform,
download=True, split=split)
test_set = Dataset(root=data_path, transform=test_transform,
split=split - 1)
return train_set, test_set
register_plugin(CelebAPlugin)
class CelebA(torchvision.datasets.ImageFolder):
url = ('https://www.dropbox.com/sh/8oqt9vytwxb3s4r/'
'AADIKlz8PR9zr6Y20qbkunrba/Img/img_align_celeba.zip?dl=1')
attr_url = ('https://www.dropbox.com/s/auexdy98c6g7y25/'
'list_attr_celeba.zip?dl=1')
filename = 'img_align_celeba.zip'
attr_filename = 'list_attr_celeba.zip'
def __init__(
self,
root,
transform=None,
if source in ('SVHN', 'STL10'):
dim_c, dim_x, dim_y = train_set[0][0].size()
uniques = np.unique(train_set.labels).tolist()
try:
uniques.remove(-1)
except ValueError:
pass
dim_l = len(uniques)
else:
dim_c, dim_x, dim_y = train_set[0][0].size()
labels = train_set.train_labels
if not isinstance(labels, list):
labels = labels.numpy()
dim_l = len(np.unique(labels))
dims = dict(x=dim_x, y=dim_y, c=dim_c, labels=dim_l)
input_names = ['images', 'targets', 'index']
self.add_dataset('train', train_set)
self.add_dataset('test', test_set)
self.set_input_names(input_names)
self.set_dims(**dims)
if scale is not None:
self.set_scale(scale)
register_plugin(TorchvisionDatasetPlugin)
def build(self, dim_z=64):
self.encoder.build(dim_out=dim_z)
self.decoder.build(dim_in=dim_z)
encoder = self.nets.encoder
decoder = self.nets.decoder
ae = AENetwork(encoder, decoder)
self.nets.ae = ae
def routine(self, inputs, targets, ae_criterion=F.mse_loss):
'''
Args:
ae_criterion: Criterion for the autoencoder.
'''
ae = self.nets.ae
outputs = ae(inputs)
r_loss = ae_criterion(outputs, inputs,
size_average=False) / inputs.size(0)
self.losses.ae = r_loss
def visualize(self, inputs, targets):
ae = self.nets.ae
outputs = ae(inputs)
self.add_image(outputs, name='reconstruction')
self.add_image(inputs, name='ground truth')
register_plugin(Autoencoder)