def setUp(self):
self.test_path = "/tmp/gan"
shutil.rmtree(self.test_path, ignore_errors=True)
os.makedirs(self.test_path, exist_ok=True)
config.gan.generator.optimizer.copy_optimizer_state = True
config.gan.generator.optimizer.type = "Adam"
config.gan.discriminator.optimizer = config.gan.generator.optimizer
config.evolution.sequential_layers = True
config.evolution.max_generations = 3
config.gan.batches_limit = 2
config.gan.batch_size = 4
config.gan.discriminator.population_size = 3
config.gan.generator.population_size = 3
config.layer.conv2d.max_channels_power = 5
config.evolution.fitness.generator = "loss"
config.evolution.fitness.fid_sample_size = 10
config.evolution.algorithm = "NEAT"
self.train = GanTrain(log_dir=self.test_path)
def store(path, size, dataset_name="CIFAR10", generator_path=None):
path = os.path.join(base_path, path)
os.makedirs(path, exist_ok=True)
if generator_path:
generator = Generator.load(generator_path)
generator_dataset = GeneratorDataset(generator, size=size)
store_from_dataset(generator_dataset, path, size)
else:
util.config.gan.dataset = dataset_name
train = GanTrain(log_dir="/tmp")
dataset = train.train_loader.dataset
store_from_dataset(dataset, path, size)
def load_data(batch_size, dataset_name, images_per_model, run_dirs):
df = pd.DataFrame()
image_shape = None
noise_data = None
for run_dir, generation in run_dirs:
target_size = len(df) + images_per_model
if generation is None:
config.gan.dataset = dataset_name
# config.gan.dataset_resize = [64, 64]
dataset = GanTrain.create_dataset()
train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, drop_last=True)
# load images from dataset
for images, labels in train_loader:
image_shape = images.shape[1:]
df_new = pd.DataFrame(images.numpy().reshape((-1, np.prod(image_shape))))
df_new["model"] = dataset_name
df_new["run_dir"] = run_dir
df_new["generation"] = None
df_new["y"] = np.zeros(len(images)) if len(labels.shape) > 1 else labels.numpy()
df = df.append(df_new)
if len(df) >= target_size:
break
else:
if noise_data is None:
noise_data = Generator().generate_noise(images_per_model, volatile=True)
print("noise data created", noise_data.shape)
last_model = sorted(glob.glob(os.path.join(run_dir, "generations", f"{generation:03}", "generator.pkl")))[
-1]
best_generator = tools.cuda(Generator.load(last_model))
n = 0
while len(df) < target_size:
noise = noise_data[n:min(n+batch_size, len(noise_data))]
n += batch_size
images = best_generator(noise).detach().cpu().numpy()
image_shape = images.shape[1:]
df_new = pd.DataFrame(images.reshape((-1, np.prod(image_shape))))
df_new["model"] = f"{run_dir}|{generation}"
df_new["run_dir"] = run_dir
df_new["generation"] = generation
df_new["y"] = np.zeros(len(images))
df = df.append(df_new)
del noise
if len(df) >= target_size:
break
best_generator = best_generator.cpu()
torch.cuda.empty_cache()
print(df.describe())
return df, image_shape
# -*- coding: future_fstrings -*-
import torch
from evolution.gan_train import GanTrain
# import better_exceptions; better_exceptions.hook()
import logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
if torch.cuda.is_available():
logger.info("CUDA device detected!")
else:
logger.info("CUDA device not detected!")
if __name__ == "__main__":
GanTrain().start()
if __name__ == "__main__":
# Test FID score using images from the input dataset
from hydra.experimental import compose, initialize
initialize(config_path="../conf", strict=True)
import util
util.config._config = compose("config.yaml")
import numpy as np
from evolution.gan_train import GanTrain
from evolution.generator import Generator, GeneratorDataset
base_path = os.path.join(os.path.dirname(__file__), "..")
fid_stat = "" #os.path.join(base_path, "fid_stats_cifar10_train.npz")
util.config.gan.dataset = "CIFAR10"
logger.info(f"start gan train")
train = GanTrain(log_dir="/tmp")
fid_size = 10000
if fid_stat and os.path.exists(fid_stat):
f = np.load(fid_stat)
m, s = f['mu'][:], f['sigma'][:]
print("fid_stats_cifar10_train", m, s)
base_fid_statistics = m, s
f.close()
inception_model = build_inception_model()
else:
logger.info(f"init fid")
initialize_fid(train.train_loader, size=fid_size)
m, s = base_fid_statistics
print("calc fid stats", m, s)
logger.info(f"finish fid")
# noise = generator.generate_noise(sample_size)
generated_images = generator(noise.cpu()).detach()
score = fid_images(generated_images)
generator.zero_grad()
return score
def fid_images(generated_images):
global base_fid_statistics, inception_model
inception_model = tools.cuda(inception_model)
m1, s1 = fid_score.calculate_activation_statistics(
generated_images.data.cpu().numpy(),
inception_model,
cuda=tools.is_cuda_available(),
dims=config.evolution.fitness.fid_dimension)
inception_model.cpu()
m2, s2 = base_fid_statistics
ret = fid_score.calculate_frechet_distance(m1, s1, m2, s2)
return ret
if __name__ == '__main__':
from evolution.gan_train import GanTrain
train = GanTrain()
generators_population, discriminators_population = train.generate_intial_population(
)
initialize_fid(train.train_loader)
score = fid(generators_population.phenotypes()[0])
logger.debug("fid", score)
def test_nsga2(self):
shutil.rmtree(self.test_path, ignore_errors=True)
config.evolution.algorithm = "NSGA2"
self.train = GanTrain(log_dir=self.test_path)
self.train.start()