import pickle
import typing
from abc import ABC, abstractmethod
from configparser import ConfigParser
import torch
from gans.causal_gan import CausalGAN
from gans.conditional_gan_cat import ConditionalCatGAN
from gans.conditional_gan_proj import ConditionalProjGAN
from gans.gan import GAN
[docs]
def parse_list(str_list: str, type_: type) -> list:
return list(map(type_, str.split(str_list)))
[docs]
class IGANFactory(ABC):
"""
Factory that represents a GAN.
This factory does not keep of created references.
"""
[docs]
def __init__(self, parser: ConfigParser) -> None:
"""
Initialize the factory.
Parameters
----------
parser : ConfigParser
Parser for config file containing GAN model and training params.
"""
self.parser = parser
[docs]
@abstractmethod
def get_gan(self) -> GAN:
"""
Returns a GAN instance
Returns
-------
GAN
GAN instance.
"""
pass
[docs]
@abstractmethod
def get_trainer(self) -> typing.Callable:
"""
Returns the GAN train function.
Returns
-------
typing.Callable
GAN train() function.
"""
pass
[docs]
class GANFactory(IGANFactory):
[docs]
def get_gan(self) -> GAN:
return GAN(
genes_no=self.parser.getint("Data", "number of genes"),
batch_size=self.parser.getint("Training", "batch size"),
latent_dim=self.parser.getint("Model", "latent dim"),
gen_layers=parse_list(self.parser["Model"]["generator layers"], int),
crit_layers=parse_list(self.parser["Model"]["critic layers"], int),
device=self.parser.get("EXPERIMENT", "device", fallback=None),
library_size=self.parser.getint("Preprocessing", "library size"),
)
[docs]
def get_trainer(self) -> typing.Callable:
gan = self.get_gan()
return lambda: gan.train(
train_files=self.parser.get("Data", "train"),
valid_files=self.parser.get("Data", "validation"),
critic_iter=self.parser.getint("Training", "critic iterations"),
max_steps=self.parser.getint("Training", "maximum steps"),
c_lambda=self.parser.getfloat("Model", "lambda"),
beta1=self.parser.getfloat("Optimizer", "beta1"),
beta2=self.parser.getfloat("Optimizer", "beta2"),
gen_alpha_0=self.parser.getfloat("Learning Rate", "generator initial"),
gen_alpha_final=self.parser.getfloat("Learning Rate", "generator final"),
crit_alpha_0=self.parser.getfloat("Learning Rate", "critic initial"),
crit_alpha_final=self.parser.getfloat("Learning Rate", "critic final"),
checkpoint=self.parser.get("EXPERIMENT", "checkpoint", fallback=None),
summary_freq=self.parser.getint("Logging", "summary frequency"),
plt_freq=self.parser.getint("Logging", "plot frequency"),
save_feq=self.parser.getint("Logging", "save frequency"),
output_dir=self.parser.get("EXPERIMENT", "output directory"),
)
[docs]
class ConditionalCatGANFactory(IGANFactory):
[docs]
def get_gan(self) -> ConditionalCatGAN:
return ConditionalCatGAN(
genes_no=self.parser.getint("Data", "number of genes"),
batch_size=self.parser.getint("Training", "batch size"),
latent_dim=self.parser.getint("Model", "latent dim"),
gen_layers=parse_list(self.parser["Model"]["generator layers"], int),
crit_layers=parse_list(self.parser["Model"]["critic layers"], int),
num_classes=self.parser.getint("Data", "number of classes"),
label_ratios=torch.Tensor(
parse_list(self.parser["Data"]["label ratios"], float)
),
device=self.parser.get("EXPERIMENT", "device", fallback=None),
library_size=self.parser.getint("Preprocessing", "library size"),
)
[docs]
def get_trainer(self) -> typing.Callable:
gan = self.get_gan()
return lambda: gan.train(
train_files=self.parser.get("Data", "train"),
valid_files=self.parser.get("Data", "validation"),
critic_iter=self.parser.getint("Training", "critic iterations"),
max_steps=self.parser.getint("Training", "maximum steps"),
c_lambda=self.parser.getfloat("Model", "lambda"),
beta1=self.parser.getfloat("Optimizer", "beta1"),
beta2=self.parser.getfloat("Optimizer", "beta2"),
gen_alpha_0=self.parser.getfloat("Learning Rate", "generator initial"),
gen_alpha_final=self.parser.getfloat("Learning Rate", "generator final"),
crit_alpha_0=self.parser.getfloat("Learning Rate", "critic initial"),
crit_alpha_final=self.parser.getfloat("Learning Rate", "critic final"),
checkpoint=self.parser.get("EXPERIMENT", "checkpoint", fallback=None),
summary_freq=self.parser.getint("Logging", "summary frequency"),
plt_freq=self.parser.getint("Logging", "plot frequency"),
save_feq=self.parser.getint("Logging", "save frequency"),
output_dir=self.parser.get("EXPERIMENT", "output directory"),
)
[docs]
class ConditionalProjGANFactory(IGANFactory):
[docs]
def get_gan(self) -> ConditionalProjGAN:
return ConditionalProjGAN(
genes_no=self.parser.getint("Data", "number of genes"),
batch_size=self.parser.getint("Training", "batch size"),
latent_dim=self.parser.getint("Model", "latent dim"),
gen_layers=parse_list(self.parser["Model"]["generator layers"], int),
crit_layers=parse_list(self.parser["Model"]["critic layers"], int),
num_classes=self.parser.getint("Data", "number of classes"),
label_ratios=torch.Tensor(
parse_list(self.parser["Data"]["label ratios"], float)
),
device=self.parser.get("EXPERIMENT", "device", fallback=None),
library_size=self.parser.getint("Preprocessing", "library size"),
)
[docs]
def get_trainer(self) -> typing.Callable:
gan = self.get_gan()
return lambda: gan.train(
train_files=self.parser.get("Data", "train"),
valid_files=self.parser.get("Data", "validation"),
critic_iter=self.parser.getint("Training", "critic iterations"),
max_steps=self.parser.getint("Training", "maximum steps"),
c_lambda=self.parser.getfloat("Model", "lambda"),
beta1=self.parser.getfloat("Optimizer", "beta1"),
beta2=self.parser.getfloat("Optimizer", "beta2"),
gen_alpha_0=self.parser.getfloat("Learning Rate", "generator initial"),
gen_alpha_final=self.parser.getfloat("Learning Rate", "generator final"),
crit_alpha_0=self.parser.getfloat("Learning Rate", "critic initial"),
crit_alpha_final=self.parser.getfloat("Learning Rate", "critic final"),
checkpoint=self.parser.get("EXPERIMENT", "checkpoint", fallback=None),
summary_freq=self.parser.getint("Logging", "summary frequency"),
plt_freq=self.parser.getint("Logging", "plot frequency"),
save_feq=self.parser.getint("Logging", "save frequency"),
output_dir=self.parser.get("EXPERIMENT", "output directory"),
)
[docs]
class CausalGANFactory(IGANFactory):
[docs]
def get_cc(self) -> GAN:
return GAN(
genes_no=self.parser.getint("Data", "number of genes"),
batch_size=self.parser.getint("CC Training", "batch size"),
latent_dim=self.parser.getint("CC Model", "latent dim"),
gen_layers=parse_list(self.parser["CC Model"]["generator layers"], int),
crit_layers=parse_list(self.parser["CC Model"]["critic layers"], int),
device=self.parser.get("EXPERIMENT", "device", fallback=None),
library_size=self.parser.getint("Preprocessing", "library size"),
)
[docs]
def get_gan(self) -> CausalGAN:
with open(self.parser.get("Data", "causal graph"), "rb") as fp:
causal_graph = pickle.load(fp)
return CausalGAN(
genes_no=self.parser.getint("Data", "number of genes"),
batch_size=self.parser.getint("Training", "batch size"),
latent_dim=self.parser.getint("Model", "latent dim"),
noise_per_gene=self.parser.getint("Model", "noise per gene"),
depth_per_gene=self.parser.getint("Model", "depth per gene"),
width_per_gene=self.parser.getint("Model", "width per gene"),
cc_latent_dim=self.parser.getint("CC Model", "latent dim"),
cc_layers=parse_list(self.parser["CC Model"]["generator layers"], int),
cc_pretrained_checkpoint=self.parser.get("EXPERIMENT", "output directory")
+ f"_CC/checkpoints/step_{self.parser.getint('CC Training', 'maximum steps')}.pth",
crit_layers=parse_list(self.parser["Model"]["critic layers"], int),
causal_graph=causal_graph,
labeler_layers=parse_list(self.parser["Model"]["labeler layers"], int),
device=self.parser.get("EXPERIMENT", "device", fallback=None),
library_size=self.parser.getint("Preprocessing", "library size"),
)
[docs]
def get_trainer(self) -> typing.Callable:
cc = self.get_cc()
# the following lambda will train the causal controller for maximum steps
# specified in the CC Training section of the config file
# after training the causal controller, the causal GAN will be instantiated
# with the pretrained causal controller and training will start.
return lambda: (
cc.train(
train_files=self.parser.get("Data", "train"),
valid_files=self.parser.get("Data", "validation"),
critic_iter=self.parser.getint("CC Training", "critic iterations"),
max_steps=self.parser.getint("CC Training", "maximum steps"),
c_lambda=self.parser.getfloat("CC Model", "lambda"),
beta1=self.parser.getfloat("CC Optimizer", "beta1"),
beta2=self.parser.getfloat("CC Optimizer", "beta2"),
gen_alpha_0=self.parser.getfloat(
"CC Learning Rate", "generator initial"
),
gen_alpha_final=self.parser.getfloat(
"CC Learning Rate", "generator final"
),
crit_alpha_0=self.parser.getfloat("CC Learning Rate", "critic initial"),
crit_alpha_final=self.parser.getfloat(
"CC Learning Rate", "critic final"
),
checkpoint=self.parser.get("EXPERIMENT", "checkpoint", fallback=None),
summary_freq=self.parser.getint("CC Logging", "summary frequency"),
plt_freq=self.parser.getint("CC Logging", "plot frequency"),
save_feq=self.parser.getint("CC Logging", "save frequency"),
output_dir=self.parser.get("EXPERIMENT", "output directory") + "_CC",
),
self.get_gan().train(
train_files=self.parser.get("Data", "train"),
valid_files=self.parser.get("Data", "validation"),
critic_iter=self.parser.getint("Training", "critic iterations"),
max_steps=self.parser.getint("Training", "maximum steps"),
c_lambda=self.parser.getfloat("Model", "lambda"),
beta1=self.parser.getfloat("Optimizer", "beta1"),
beta2=self.parser.getfloat("Optimizer", "beta2"),
gen_alpha_0=self.parser.getfloat("Learning Rate", "generator initial"),
gen_alpha_final=self.parser.getfloat(
"Learning Rate", "generator final"
),
crit_alpha_0=self.parser.getfloat("Learning Rate", "critic initial"),
crit_alpha_final=self.parser.getfloat("Learning Rate", "critic final"),
labeler_alpha=self.parser.getfloat("Learning Rate", "labeler"),
antilabeler_alpha=self.parser.getfloat("Learning Rate", "antilabeler"),
labeler_training_interval=self.parser.getfloat(
"Training", "labeler and antilabeler training intervals"
),
checkpoint=self.parser.get("EXPERIMENT", "checkpoint", fallback=None),
summary_freq=self.parser.getint("Logging", "summary frequency"),
plt_freq=self.parser.getint("Logging", "plot frequency"),
save_feq=self.parser.getint("Logging", "save frequency"),
output_dir=self.parser.get("EXPERIMENT", "output directory"),
),
)[0]
[docs]
def get_factory(cfg: ConfigParser) -> IGANFactory:
"""
Return the factory for the GAN type based on 'model' key in the parser.
Parameters
----------
cfg : ConfigParser
Parser for config file containing GAN model and training params.
Returns
-------
IGANFactory
Factory for the specified GAN.
Raises
------
ValueError
If the model is unknown or not implemented.
"""
# read the desired GAN
model = cfg.get("Model", "type")
factories = {
"GAN": GANFactory(cfg),
"proj conditional GAN": ConditionalProjGANFactory(cfg),
"cat conditional GAN": ConditionalCatGANFactory(cfg),
"causal GAN": CausalGANFactory(cfg),
}
if model in factories:
return factories[model]
raise ValueError(f"model '{model}' type is invalid")
