def train_generator_step(self, configuration: Configuration, metadata: Metadata,
architecture: Architecture) -> float:
# clean previous gradients
architecture.generator_optimizer.zero_grad()
# conditional
if "conditional" in architecture.arguments:
# for now uniform distribution is used but could be controlled in a different way
# also this works for both binary and categorical dependent variables
number_of_conditions = metadata.get_dependent_variable().get_size()
condition = to_gpu_if_available(FloatTensor(configuration.batch_size).uniform_(0, number_of_conditions))
# non-conditional
else:
condition = None
# generate a full batch of fake features
fake_features = self.sample_fake(architecture, configuration.batch_size, condition=condition)
# calculate loss
loss = architecture.generator_loss(architecture, fake_features, condition=condition)
# calculate gradients
loss.backward()
# update the generator weights
architecture.generator_optimizer.step()
# return the loss
return to_cpu_if_was_in_gpu(loss).item()
def run(self, configuration: Configuration) -> None:
seed_all(configuration.get("seed"))
metadata = load_metadata(configuration.metadata)
architecture_configuration = load_configuration(configuration.architecture)
self.validate_architecture_configuration(architecture_configuration)
architecture = create_architecture(metadata, architecture_configuration)
architecture.to_gpu_if_available()
checkpoints = Checkpoints()
checkpoint = checkpoints.load(configuration.checkpoint)
if "best_architecture" in checkpoint:
checkpoints.load_states(checkpoint["best_architecture"], architecture)
else:
checkpoints.load_states(checkpoint["architecture"], architecture)
# pre-processing
imputation = create_component(architecture, metadata, configuration.imputation)
pre_processing = PreProcessing(imputation)
# post-processing
if "scale_transform" in configuration:
scale_transform = load_scale_transform(configuration.scale_transform)
else:
scale_transform = None
post_processing = PostProcessing(metadata, scale_transform)
# load the features
features = to_gpu_if_available(torch.from_numpy(np.load(configuration.features)).float())
missing_mask = to_gpu_if_available(torch.from_numpy(np.load(configuration.missing_mask)).float())
# initial imputation
batch = pre_processing.transform({"features": features, "missing_mask": missing_mask})
# generate the model outputs
output = self.impute(configuration, metadata, architecture, batch)
# imputation
output = compose_with_mask(mask=missing_mask, differentiable=False, where_one=output, where_zero=features)
# post-process
output = post_processing.transform(output)
# save the imputation
output = to_cpu_if_was_in_gpu(output)
output = output.numpy()
np.save(configuration.output, output)
def run(self, configuration: Configuration) -> None:
seed_all(configuration.get("seed"))
metadata = load_metadata(configuration.metadata)
architecture_configuration = load_configuration(
configuration.architecture)
self.validate_architecture_configuration(architecture_configuration)
architecture = create_architecture(metadata,
architecture_configuration)
architecture.to_gpu_if_available()
checkpoints = Checkpoints()
checkpoint = checkpoints.load(configuration.checkpoint)
if "best_architecture" in checkpoint:
checkpoints.load_states(checkpoint["best_architecture"],
architecture)
else:
checkpoints.load_states(checkpoint["architecture"], architecture)
# load the features
features = to_gpu_if_available(
torch.from_numpy(np.load(configuration.features)).float())
# conditional
if "labels" in configuration:
condition = to_gpu_if_available(
torch.from_numpy(np.load(configuration.labels)).float())
else:
condition = None
# encode
with torch.no_grad():
code = architecture.autoencoder.encode(features,
condition=condition)["code"]
# save the code
code = to_cpu_if_was_in_gpu(code)
code = code.numpy()
np.save(configuration.output, code)
def train_discriminator_step(self, configuration: Configuration, metadata: Metadata, architecture: Architecture,
batch: Batch) -> float:
# clean previous gradients
architecture.discriminator_optimizer.zero_grad()
# generate a batch of fake features with the same size as the real feature batch
fake_features = self.sample_fake(architecture, len(batch["features"]), condition=batch.get("labels"))
fake_features = fake_features.detach() # do not propagate to the generator
# calculate loss
loss = architecture.discriminator_loss(architecture,
batch["features"],
fake_features,
condition=batch.get("labels"))
# calculate gradients
loss.backward()
# update the discriminator weights
architecture.discriminator_optimizer.step()
# return the loss
return to_cpu_if_was_in_gpu(loss).item()
def train_generator_step(configuration: Configuration, metadata: Metadata,
architecture: Architecture,
batch: Batch) -> float:
# clean previous gradients
architecture.generator_optimizer.zero_grad()
# generate a batch of fake features with the same size as the real feature batch
generated = architecture.generator(batch["features"],
missing_mask=batch["missing_mask"])
# replace the missing features by the generated
imputed = compose_with_mask(
mask=batch["missing_mask"],
differentiable=True, # now there are no NaNs and this should be used
where_one=generated,
where_zero=batch["raw_features"])
# generate hint
hint = generate_hint(batch["missing_mask"],
configuration.hint_probability, metadata)
# calculate loss
loss = architecture.generator_loss(architecture=architecture,
features=batch["raw_features"],
generated=generated,
imputed=imputed,
hint=hint,
non_missing_mask=inverse_mask(
batch["missing_mask"]))
# calculate gradients
loss.backward()
# update the generator weights
architecture.generator_optimizer.step()
# return the loss
return to_cpu_if_was_in_gpu(loss).item()
def run(self, configuration: Configuration) -> None:
seed_all(configuration.get("seed"))
metadata = load_metadata(configuration.metadata)
if "scale_transform" in configuration:
scale_transform = load_scale_transform(
configuration.scale_transform)
else:
scale_transform = None
post_processing = PostProcessing(metadata, scale_transform)
architecture_configuration = load_configuration(
configuration.architecture)
self.validate_architecture_configuration(architecture_configuration)
architecture = create_architecture(metadata,
architecture_configuration)
architecture.to_gpu_if_available()
checkpoints = Checkpoints()
checkpoint = checkpoints.load(configuration.checkpoint)
if "best_architecture" in checkpoint:
checkpoints.load_states(checkpoint["best_architecture"],
architecture)
else:
checkpoints.load_states(checkpoint["architecture"], architecture)
samples = []
# create the strategy if defined
if "strategy" in configuration:
# validate strategy name is present
if "factory" not in configuration.strategy:
raise Exception(
"Missing factory name while creating sample strategy.")
# validate strategy name
strategy_name = configuration.strategy.factory
if strategy_name not in strategy_class_by_name:
raise Exception(
"Invalid factory name '{}' while creating sample strategy."
.format(strategy_name))
# create the strategy
strategy_class = strategy_class_by_name[strategy_name]
strategy = strategy_class(**configuration.strategy.get(
"arguments", default={}, transform_default=False))
# use the default strategy
else:
strategy = DefaultSampleStrategy()
# this is only to pass less parameters back and forth
sampler = Sampler(self, configuration, metadata, architecture,
post_processing)
# while more samples are needed
start = 0
while start < configuration.sample_size:
# do not calculate gradients
with torch.no_grad():
# sample:
# the task delegates to the strategy and passes the sampler object to avoid passing even more parameters
# the strategy may prepare additional sampling arguments (e.g. condition)
# the strategy delegates to the sampler object
# the sampler object delegates back to the task adding parameters that it was keeping
# the task child class does the actual sampling depending on the model
# the sampler object applies post-processing
# the strategy may apply filtering to the samples (e.g. rejection)
# the task finally gets the sample
batch_samples = strategy.generate_sample(
sampler, configuration, metadata)
# transform back the samples
batch_samples = to_cpu_if_was_in_gpu(batch_samples)
batch_samples = batch_samples.numpy()
# if the batch is not empty
if len(batch_samples) > 0:
# do not go further than the desired number of samples
end = min(start + len(batch_samples),
configuration.sample_size)
# limit the samples taken from the batch based on what is missing
batch_samples = batch_samples[:min(len(batch_samples), end -
start), :]
# if it is the first batch
if len(samples) == 0:
samples = batch_samples
# if its not the first batch we have to concatenate
else:
samples = np.concatenate((samples, batch_samples), axis=0)
# move to next batch
start = end
# save the samples
np.save(configuration.output, samples)
def impute(self, configuration: Configuration, metadata: Metadata,
architecture: Architecture, batch: Dict[str, Tensor]) -> Tensor:
# loss function
loss_function = create_component(architecture, metadata,
configuration.reconstruction_loss)
masked_loss_function = MaskedReconstructionLoss(loss_function)
batch_size = batch["features"].shape[0] * batch["features"].shape[1]
# we need the non missing mask for the loss
non_missing_mask = inverse_mask(batch["missing_mask"])
# initial noise
noise = to_gpu_if_available(
FloatTensor(len(batch["features"]),
architecture.arguments.noise_size).normal_())
noise.requires_grad_()
# it is not the generator what we are updating
# it is the noise
optimizer = Adam([noise],
weight_decay=0,
lr=configuration.noise_learning_rate)
architecture.generator.eval()
# logger
log_path = create_parent_directories_if_needed(configuration.logs)
logger = TrainLogger(self.logger, log_path, False)
# initial generation
logger.start_timer()
generated = architecture.generator(noise,
condition=batch.get("labels"))
# iterate until we reach the maximum number of iterations or until the non missing loss is too small
max_iterations = configuration.max_iterations
for iteration in range(1, max_iterations + 1):
# compute the loss on the non-missing values
non_missing_loss = masked_loss_function(generated,
batch["features"],
non_missing_mask)
logger.log(iteration, max_iterations, "non_missing_loss",
to_cpu_if_was_in_gpu(non_missing_loss).item())
# this loss only makes sense if the ground truth is present
# only used for debugging
if configuration.get("log_missing_loss", False):
# this part should not affect the gradient calculation
with torch.no_grad():
missing_loss = masked_loss_function(
generated, batch["raw_features"],
batch["missing_mask"])
logger.log(iteration, max_iterations, "missing_loss",
to_cpu_if_was_in_gpu(missing_loss).item())
loss = loss_function(generated,
batch["raw_features"]) / batch_size
logger.log(iteration, max_iterations, "loss",
to_cpu_if_was_in_gpu(loss).item())
# if the generation is good enough we stop
if to_cpu_if_was_in_gpu(non_missing_loss).item(
) < configuration.get("tolerance", 1e-5):
break
# clear previous gradients
optimizer.zero_grad()
# compute the gradients
non_missing_loss.backward()
# update the noise
optimizer.step()
# generate next
logger.start_timer()
generated = architecture.generator(noise,
condition=batch.get("labels"))
return generated
def val_batch(architecture: Architecture, batch: Batch,
post_processing: PostProcessing) -> float:
generated = architecture.generator(batch["features"],
missing_mask=batch["missing_mask"])
loss = architecture.val_loss(post_processing, generated, batch)
return to_cpu_if_was_in_gpu(loss).item()
def to_cpu_if_was_in_gpu(self) -> None:
for name, component in self.items():
if isinstance(component, Module): # skip optimizers
self[name] = to_cpu_if_was_in_gpu(component)