def evaluate_batch(images_train: torch.Tensor,
labels_train: torch.Tensor,
images_valid: torch.Tensor,
labels_valid: torch.Tensor,
plot_images=False,
iteration: int = 0) -> dict:
""" evaluate BigGAN model """
generator.eval()
discriminator.eval()
latent = generate_latent(labels_train.shape[0], latent_dim, sigma=False)
# labels_fake = generate_labels(labels_train.shape[0])
labels_fake = labels_valid # generated images should fit the distribution of validation data. both are checked against training data
generated_images = generator(latent, labels_fake)
if plot_images:
width = min(8, len(generated_images))
write_generated_galaxy_images_iteration(iteration=iteration, images=generated_images.cpu(), width=width, height=len(generated_images)//width)
write_generated_galaxy_images_iteration(iteration=0, images=images_train.cpu(), width=width, height=len(generated_images)//width, file_prefix="original_sample")
d_loss_real_train, d_loss_fake, accuracy_real_train, accuracy_fake = compute_loss_discriminator(images_train, labels_train, generated_images, labels_fake, accuracy=True)
d_loss_real_valid, d_loss_fake, accuracy_real_valid, accuracy_fake = compute_loss_discriminator(images_valid, labels_valid, generated_images, labels_fake, accuracy=True)
g_loss_dis, g_loss_class, accuracy_gen_dis, accuracy_gen_class = compute_loss_generator(generated_images, labels_fake, accuracy=True)
# variance_train = get_sample_variance(images_train)
# variance_generator = get_sample_variance(generated_images)
# divergence_intensity_valid = compute_divergence_total_intensity_statistics(images_train, images_valid)
# divergence_intensity_generator = compute_divergence_total_intensity_statistics(images_train, generated_images)
# divergence_bluriness_valid = compute_divergence_bluriness_statistics(images_train, images_valid)
# divergence_bluriness_generator = compute_divergence_bluriness_statistics(images_train, generated_images)
logs = {
"loss_discriminator_real_train": d_loss_real_train.item(),
"loss_discriminator_real_valid": d_loss_real_valid.item(),
"loss_discriminator_fake": d_loss_fake.item(),
"loss_generator_dis": g_loss_dis.item(),
### "variance_train" : variance_train.item(),
### "variance_generator" : variance_generator.item(),
# "divergence_intensity_valid" : divergence_intensity_valid.item(),
# "divergence_intensity_generator": divergence_intensity_generator.item(),
### "divergence_bluriness_valid": divergence_bluriness_valid.item(),
### "divergence_bluriness_generator": divergence_bluriness_generator.item(),
"accuracy_real_train" : accuracy_real_train.item(),
"accuracy_real_valid": accuracy_real_valid.item(),
"accuracy_fake": accuracy_fake.item(),
"accuracy_generator_dis": accuracy_gen_dis.item(),
}
if conditional:
logs.update({
"loss_generator_class" : g_loss_class.item(),
"accuracy_generator_class": accuracy_gen_class.item(),
})
return logs
def test_model(model):
# load labels
make_data_loader = MakeDataLoader(N_sample=64, augmented=False)
data_loader_valid = make_data_loader.get_data_loader_valid(batch_size=N_sample, shuffle=False)
for images_truth, labels in data_loader_valid:
break
# generate latent
latent = generate_latent(labels.shape[0], model.dim_z, sigma=False, device="cpu")
# create images
images = model(latent, labels)
# plot images
write_generated_galaxy_images_iteration(iteration=0, images=images_truth.detach().cpu(), width=8, height=8, file_prefix="original_sample")
write_generated_galaxy_images_iteration(iteration=1, images=images.detach().cpu(), width=8, height=8, file_prefix=f"test_generator_{type(model).__name__}_sample")
def train_generator(batch_size: int = mini_batch_size,
optimizer_step: bool = True):
latent = generate_latent(batch_size, latent_dim, sigma=False)
labels_fake = generate_labels(batch_size)
if not conditional:
labels_fake[:] = 0
generated_images = generator(latent, labels_fake)
g_loss_dis, g_loss_class = compute_loss_generator(generated_images, labels_fake)
g_loss = g_loss_dis
if conditional:
g_loss += g_loss_class
g_loss.backward()
if optimizer_step:
generator.optimizer.step()
generator.zero_grad()
def train_discriminator(images: torch.Tensor,
labels: torch.Tensor,
optimizer_step: bool = True):
latent = generate_latent(labels.shape[0], latent_dim, sigma=False)
labels_fake = generate_labels(labels.shape[0])
if not conditional:
labels_fake[:] = 0
generated_images = generator(latent, labels_fake)
d_loss_real, d_loss_fake = compute_loss_discriminator(images, labels, generated_images, labels_fake)
d_loss = d_loss_real + d_loss_fake
d_loss.backward()
if optimizer_step:
discriminator.optimizer.step()
discriminator.zero_grad()
def train_step(images: torch.Tensor,
labels: torch.Tensor,
train_generator: bool = True):
latent = generate_latent(labels.shape[0], latent_dim, sigma=False)
# for generator training, use 50 % generated labels, 50 % training labels with noise
labels_fake = generate_labels(labels.shape[0])
half = labels.shape[0] // 2
labels_fake[:half] = labels[:half] + noise_label * torch.randn(
labels[:half].shape).to(device)
labels_fake = make_galaxy_labels_hierarchical(labels_fake)
# for discriminator training, add noise on input labels
labels += noise_label * torch.randn(labels.shape).to(device)
labels = make_galaxy_labels_hierarchical(labels)
generated_images = generator(latent, labels_fake)
# disc training
d_loss_real, d_loss_fake = compute_loss_discriminator(
images, labels, generated_images, labels_fake)
d_loss = 0.5 * torch.add(d_loss_real, d_loss_fake)
discriminator.zero_grad()
d_loss.backward()
discriminator.optim.step()
# gen training
if train_generator:
g_loss_dis, g_loss_class = compute_loss_generator(
generated_images, labels_fake)
g_loss = g_loss_dis
if conditional:
g_loss += g_loss_class
generator.zero_grad()
g_loss.backward()
generator.optim.step()