def __init__(self, n_input: int, n_batch: int = 0, n_labels: int = 0,
n_hidden: int = 128, n_latent: int = 10, n_layers: int = 1,
dropout_rate: float = 0.1, dispersion: str = "gene",
log_variational: bool = True, reconstruction_loss: str = "zinb",
y_prior=None, labels_groups: Sequence[int] = None, use_labels_groups: bool = False,
classifier_parameters: dict = dict()):
super().__init__(n_input, n_hidden=n_hidden, n_latent=n_latent, n_layers=n_layers,
dropout_rate=dropout_rate, n_batch=n_batch, dispersion=dispersion,
log_variational=log_variational, reconstruction_loss=reconstruction_loss)
self.n_labels = n_labels
# Classifier takes n_latent as input
cls_parameters = {"n_layers": n_layers, "n_hidden": n_hidden, "dropout_rate": dropout_rate}
cls_parameters.update(classifier_parameters)
self.classifier = Classifier(n_latent, n_labels=n_labels, **cls_parameters)
self.encoder_z2_z1 = Encoder(n_latent, n_latent, n_cat_list=[self.n_labels], n_layers=n_layers,
n_hidden=n_hidden, dropout_rate=dropout_rate)
self.decoder_z1_z2 = Decoder(n_latent, n_latent, n_cat_list=[self.n_labels], n_layers=n_layers,
n_hidden=n_hidden)
self.y_prior = torch.nn.Parameter(
y_prior if y_prior is not None else (1 / n_labels) * torch.ones(1, n_labels), requires_grad=False
)
self.use_labels_groups = use_labels_groups
self.labels_groups = np.array(labels_groups) if labels_groups is not None else None
if self.use_labels_groups:
assert labels_groups is not None, "Specify label groups"
unique_groups = np.unique(self.labels_groups)
self.n_groups = len(unique_groups)
assert (unique_groups == np.arange(self.n_groups)).all()
self.classifier_groups = Classifier(n_latent, n_hidden, self.n_groups, n_layers, dropout_rate)
self.groups_index = torch.nn.ParameterList([torch.nn.Parameter(
torch.tensor((self.labels_groups == i).astype(np.uint8), dtype=torch.uint8), requires_grad=False
) for i in range(self.n_groups)])
def train(self, n_epochs=20, lr=1e-3, weight_decay=1e-6, params=None):
self.adversarial_cls = Classifier(self.model.n_latent, n_labels=self.model.n_batch, n_layers=3)
if self.use_cuda:
self.adversarial_cls.cuda()
self.optimizer_cls = torch.optim.Adam(filter(lambda p: p.requires_grad, self.adversarial_cls.parameters()),
lr=lr, weight_decay=weight_decay)
super(AdversarialTrainerVAE, self).train(n_epochs=n_epochs, lr=lr, params=None)
class AdversarialTrainerVAE(Trainer):
r"""The modified UnsupervisedTrainer class for the unsupervised training of an autoencoder.
"""
default_metrics_to_monitor = ['ll']
def __init__(self, model, gene_dataset, train_size=0.8, test_size=None,
n_epochs_even=1, n_epochs_cl=1, warm_up=10,
scale=50, **kwargs):
super(AdversarialTrainerVAE, self).__init__(model, gene_dataset, **kwargs)
print("I am the adversarial Trainer")
self.kl = None
self.n_epochs_cl = n_epochs_cl
self.n_epochs_even = n_epochs_even
self.weighting = 0
self.kl_weight = 0
self.classification_ponderation = 0
self.warm_up = warm_up
self.scale = scale
self.train_set, self.test_set = self.train_test(model, gene_dataset, train_size, test_size)
self.train_set.to_monitor = ['ll']
self.test_set.to_monitor = ['ll']
@property
def posteriors_loop(self):
return ['train_set']
def train(self, n_epochs=20, lr=1e-3, weight_decay=1e-6, params=None):
self.adversarial_cls = Classifier(self.model.n_latent, n_labels=self.model.n_batch, n_layers=3)
if self.use_cuda:
self.adversarial_cls.cuda()
self.optimizer_cls = torch.optim.Adam(filter(lambda p: p.requires_grad, self.adversarial_cls.parameters()),
lr=lr, weight_decay=weight_decay)
super(AdversarialTrainerVAE, self).train(n_epochs=n_epochs, lr=lr, params=None)
def loss(self, tensors):
sample_batch, local_l_mean, local_l_var, batch_index, _ = tensors
reconst_loss, kl_divergence = self.model(sample_batch, local_l_mean, local_l_var, batch_index)
loss = torch.mean(reconst_loss + self.kl_weight * kl_divergence)
if self.epoch > self.warm_up:
z = self.model.sample_from_posterior_z(sample_batch)
cls_loss = (self.scale * F.cross_entropy(self.adversarial_cls(z), torch.zeros_like(batch_index).view(-1)))
self.optimizer_cls.zero_grad()
cls_loss.backward(retain_graph=True)
self.optimizer_cls.step()
else:
cls_loss = 0
return loss - cls_loss
def on_epoch_begin(self):
self.kl_weight = self.kl if self.kl is not None else min(1, self.epoch / 400)
def __init__(self, n_input, indexes_fish_train=None, n_batch=0, n_labels=0, n_hidden=128, n_latent=10,
n_layers=1, n_layers_decoder=1, dropout_rate=0.3,
dispersion="gene", log_variational=True, reconstruction_loss="zinb",
reconstruction_loss_fish="poisson", model_library=False):
super().__init__(n_input, dispersion=dispersion, n_latent=n_hidden, n_hidden=n_hidden,
log_variational=log_variational, dropout_rate=dropout_rate, n_layers=1,
reconstruction_loss=reconstruction_loss, n_batch=n_batch, n_labels=n_labels)
self.n_input = n_input
self.n_input_fish = len(indexes_fish_train)
self.indexes_to_keep = indexes_fish_train
self.reconstruction_loss_fish = reconstruction_loss_fish
self.model_library = model_library
self.n_latent = n_latent
# First layer of the encoder isn't shared
self.z_encoder_fish = Encoder(self.n_input_fish, n_hidden, n_hidden=n_hidden, n_layers=1,
dropout_rate=dropout_rate)
# The last layers of the encoder are shared
self.z_final_encoder = Encoder(n_hidden, n_latent, n_hidden=n_hidden, n_layers=n_layers,
dropout_rate=dropout_rate)
self.l_encoder_fish = Encoder(self.n_input_fish, 1, n_hidden=n_hidden, n_layers=1,
dropout_rate=dropout_rate)
self.l_encoder = Encoder(n_input, 1, n_hidden=n_hidden, n_layers=1,
dropout_rate=dropout_rate)
self.decoder = DecoderSCVI(n_latent, n_input, n_layers=n_layers_decoder, n_hidden=n_hidden,
n_cat_list=[n_batch])
self.classifier = Classifier(n_latent, n_labels=n_labels, n_hidden=128, n_layers=3)
def __init__(
self,
n_input,
n_batch,
n_labels,
n_hidden=128,
n_latent=10,
n_layers=1,
dropout_rate=0.1,
y_prior=None,
dispersion="gene",
log_variational=True,
reconstruction_loss="zinb",
):
super().__init__(
n_input,
n_batch,
n_labels,
n_hidden=n_hidden,
n_latent=n_latent,
n_layers=n_layers,
dropout_rate=dropout_rate,
dispersion=dispersion,
log_variational=log_variational,
reconstruction_loss=reconstruction_loss,
)
self.z_encoder = Encoder(
n_input,
n_latent,
n_cat_list=[n_labels],
n_hidden=n_hidden,
n_layers=n_layers,
dropout_rate=dropout_rate,
)
self.decoder = DecoderSCVI(
n_latent,
n_input,
n_cat_list=[n_batch, n_labels],
n_layers=n_layers,
n_hidden=n_hidden,
)
self.y_prior = torch.nn.Parameter(
y_prior if y_prior is not None else
(1 / n_labels) * torch.ones(1, n_labels),
requires_grad=False,
)
self.classifier = Classifier(n_input,
n_hidden,
n_labels,
n_layers=n_layers,
dropout_rate=dropout_rate)
def test_classifier_accuracy(save_path):
cortex_dataset = CortexDataset(save_path=save_path)
cls = Classifier(cortex_dataset.nb_genes, n_labels=cortex_dataset.n_labels)
cls_trainer = ClassifierTrainer(cls,
cortex_dataset,
metrics_to_monitor=['accuracy'],
frequency=1,
early_stopping_kwargs={
'early_stopping_metric': 'accuracy',
'save_best_state_metric': 'accuracy'
})
cls_trainer.train(n_epochs=2)
cls_trainer.train_set.accuracy()
def test_sampling_zl(save_path):
cortex_dataset = CortexDataset(save_path=save_path)
cortex_vae = VAE(cortex_dataset.nb_genes, cortex_dataset.n_batches)
trainer_cortex_vae = UnsupervisedTrainer(
cortex_vae, cortex_dataset, train_size=0.5, use_cuda=use_cuda
)
trainer_cortex_vae.train(n_epochs=2)
cortex_cls = Classifier((cortex_vae.n_latent + 1), n_labels=cortex_dataset.n_labels)
trainer_cortex_cls = ClassifierTrainer(
cortex_cls, cortex_dataset, sampling_model=cortex_vae, sampling_zl=True
)
trainer_cortex_cls.train(n_epochs=2)
trainer_cortex_cls.test_set.accuracy()
def __init__(self,
n_input,
n_batch,
n_labels,
n_hidden=128,
n_latent=10,
n_layers=1,
dropout_rate=0.1,
y_prior=None,
logreg_classifier=False,
dispersion="gene",
log_variational=True,
reconstruction_loss="zinb"):
super(SVAEC, self).__init__(n_input,
n_hidden=n_hidden,
n_latent=n_latent,
n_layers=n_layers,
dropout_rate=dropout_rate,
n_batch=n_batch,
dispersion=dispersion,
log_variational=log_variational,
reconstruction_loss=reconstruction_loss)
self.n_labels = n_labels
self.n_latent_layers = 2
# Classifier takes n_latent as input
if logreg_classifier:
self.classifier = LinearLogRegClassifier(n_latent, self.n_labels)
else:
self.classifier = Classifier(n_latent, n_hidden, self.n_labels,
n_layers, dropout_rate)
self.encoder_z2_z1 = Encoder(n_latent,
n_latent,
n_cat_list=[self.n_labels],
n_layers=n_layers,
n_hidden=n_hidden,
dropout_rate=dropout_rate)
self.decoder_z1_z2 = Decoder(n_latent,
n_latent,
n_cat_list=[self.n_labels],
n_layers=n_layers,
n_hidden=n_hidden,
dropout_rate=dropout_rate)
self.y_prior = torch.nn.Parameter(y_prior if y_prior is not None else
(1 / self.n_labels) *
torch.ones(self.n_labels),
requires_grad=False)
def test_cortex():
cortex_dataset = CortexDataset()
vae = VAE(cortex_dataset.nb_genes, cortex_dataset.n_batches)
infer_cortex_vae = VariationalInference(vae,
cortex_dataset,
train_size=0.1,
use_cuda=use_cuda)
infer_cortex_vae.train(n_epochs=1)
infer_cortex_vae.ll('train')
infer_cortex_vae.differential_expression_stats('train')
infer_cortex_vae.differential_expression('test')
infer_cortex_vae.imputation('train', corruption='uniform')
infer_cortex_vae.imputation('test', n_samples=2, corruption='binomial')
svaec = SVAEC(cortex_dataset.nb_genes, cortex_dataset.n_batches,
cortex_dataset.n_labels)
infer_cortex_svaec = JointSemiSupervisedVariationalInference(
svaec,
cortex_dataset,
n_labelled_samples_per_class=50,
use_cuda=use_cuda)
infer_cortex_svaec.train(n_epochs=1)
infer_cortex_svaec.accuracy('labelled')
infer_cortex_svaec.ll('all')
svaec = SVAEC(cortex_dataset.nb_genes,
cortex_dataset.n_batches,
cortex_dataset.n_labels,
logreg_classifier=True)
infer_cortex_svaec = AlternateSemiSupervisedVariationalInference(
svaec,
cortex_dataset,
n_labelled_samples_per_class=50,
use_cuda=use_cuda)
infer_cortex_svaec.train(n_epochs=1, lr=1e-2)
infer_cortex_svaec.accuracy('unlabelled')
infer_cortex_svaec.svc_rf(unit_test=True)
cls = Classifier(cortex_dataset.nb_genes, n_labels=cortex_dataset.n_labels)
infer_cls = ClassifierInference(cls, cortex_dataset)
infer_cls.train(n_epochs=1)
infer_cls.accuracy('train')
def test_cortex(save_path):
cortex_dataset = CortexDataset(save_path=save_path)
vae = VAE(cortex_dataset.nb_genes, cortex_dataset.n_batches)
trainer_cortex_vae = UnsupervisedTrainer(vae, cortex_dataset, train_size=0.5, use_cuda=use_cuda)
trainer_cortex_vae.train(n_epochs=1)
trainer_cortex_vae.train_set.ll()
trainer_cortex_vae.train_set.differential_expression_stats()
trainer_cortex_vae.corrupt_posteriors(corruption='binomial')
trainer_cortex_vae.corrupt_posteriors()
trainer_cortex_vae.train(n_epochs=1)
trainer_cortex_vae.uncorrupt_posteriors()
trainer_cortex_vae.train_set.imputation_benchmark(n_samples=1, show_plot=False,
title_plot='imputation', save_path=save_path)
svaec = SCANVI(cortex_dataset.nb_genes, cortex_dataset.n_batches, cortex_dataset.n_labels)
trainer_cortex_svaec = JointSemiSupervisedTrainer(svaec, cortex_dataset,
n_labelled_samples_per_class=3,
use_cuda=use_cuda)
trainer_cortex_svaec.train(n_epochs=1)
trainer_cortex_svaec.labelled_set.accuracy()
trainer_cortex_svaec.full_dataset.ll()
svaec = SCANVI(cortex_dataset.nb_genes, cortex_dataset.n_batches, cortex_dataset.n_labels)
trainer_cortex_svaec = AlternateSemiSupervisedTrainer(svaec, cortex_dataset,
n_labelled_samples_per_class=3,
use_cuda=use_cuda)
trainer_cortex_svaec.train(n_epochs=1, lr=1e-2)
trainer_cortex_svaec.unlabelled_set.accuracy()
data_train, labels_train = trainer_cortex_svaec.labelled_set.raw_data()
data_test, labels_test = trainer_cortex_svaec.unlabelled_set.raw_data()
compute_accuracy_svc(data_train, labels_train, data_test, labels_test,
param_grid=[{'C': [1], 'kernel': ['linear']}])
compute_accuracy_rf(data_train, labels_train, data_test, labels_test,
param_grid=[{'max_depth': [3], 'n_estimators': [10]}])
cls = Classifier(cortex_dataset.nb_genes, n_labels=cortex_dataset.n_labels)
cls_trainer = ClassifierTrainer(cls, cortex_dataset)
cls_trainer.train(n_epochs=1)
cls_trainer.train_set.accuracy()
def __init__(self,
n_input,
n_batch,
n_labels,
n_hidden=128,
n_latent=10,
n_layers=1,
dropout_rate=0.1,
y_prior=None,
logreg_classifier=False,
dispersion="gene",
log_variational=True,
reconstruction_loss="zinb",
labels_groups=None,
use_labels_groups=False):
super(SVAEC, self).__init__(n_input,
n_hidden=n_hidden,
n_latent=n_latent,
n_layers=n_layers,
dropout_rate=dropout_rate,
n_batch=n_batch,
dispersion=dispersion,
log_variational=log_variational,
reconstruction_loss=reconstruction_loss)
self.n_labels = n_labels
self.n_latent_layers = 2
# Classifier takes n_latent as input
if logreg_classifier:
self.classifier = LinearLogRegClassifier(n_latent, self.n_labels)
else:
self.classifier = Classifier(n_latent, n_hidden, self.n_labels,
n_layers, dropout_rate)
self.encoder_z2_z1 = Encoder(n_latent,
n_latent,
n_cat_list=[self.n_labels],
n_layers=n_layers,
n_hidden=n_hidden,
dropout_rate=dropout_rate)
self.decoder_z1_z2 = Decoder(n_latent,
n_latent,
n_cat_list=[self.n_labels],
n_layers=n_layers,
n_hidden=n_hidden,
dropout_rate=dropout_rate)
self.y_prior = torch.nn.Parameter(y_prior if y_prior is not None else
(1 / n_labels) *
torch.ones(1, n_labels),
requires_grad=False)
self.use_labels_groups = use_labels_groups
self.labels_groups = np.array(
labels_groups) if labels_groups is not None else None
if self.use_labels_groups:
assert labels_groups is not None, "Specify label groups"
unique_groups = np.unique(self.labels_groups)
self.n_groups = len(unique_groups)
assert (unique_groups == np.arange(self.n_groups)).all()
self.classifier_groups = Classifier(n_latent, n_hidden,
self.n_groups, n_layers,
dropout_rate)
self.groups_index = torch.nn.ParameterList([
torch.nn.Parameter(torch.tensor(
(self.labels_groups == i).astype(np.uint8),
dtype=torch.uint8),
requires_grad=False)
for i in range(self.n_groups)
])
class TrainerFish(Trainer):
r"""The VariationalInference class for the unsupervised training of an autoencoder.
Args:
:model: A model instance from class ``VAEF``
:gene_dataset: A gene_dataset instance like ``CortexDataset()``
:train_size: The train size, either a float between 0 and 1 or and integer for the number of training samples
to use Default: ``0.8``.
:\*\*kwargs: Other keywords arguments from the general Trainer class.
Examples:
>>> gene_dataset_seq = CortexDataset()
>>> gene_dataset_fish = SmfishDataset()
>>> vaef = VAEF(gene_dataset_seq.nb_genes, gene_dataset_fish.nb_genes,
... n_labels=gene_dataset.n_labels, use_cuda=True)
>>> trainer = TrainerFish(gene_dataset_seq, gene_dataset_fish, vaef, train_size=0.5)
>>> trainer.train(n_epochs=20, lr=1e-3)
"""
default_metrics_to_monitor = ["reconstruction_error"]
def __init__(self,
model,
gene_dataset_seq,
gene_dataset_fish,
train_size=0.8,
test_size=None,
use_cuda=True,
cl_ratio=0,
n_epochs_even=1,
n_epochs_kl=2000,
n_epochs_cl=1,
seed=0,
warm_up=10,
scale=50,
**kwargs):
super().__init__(model, gene_dataset_seq, use_cuda=use_cuda, **kwargs)
self.kl = None
self.cl_ratio = cl_ratio
self.n_epochs_cl = n_epochs_cl
self.n_epochs_even = n_epochs_even
self.n_epochs_kl = n_epochs_kl
self.weighting = 0
self.kl_weight = 0
self.classification_ponderation = 0
self.warm_up = warm_up
self.scale = scale
self.train_seq, self.test_seq = self.train_test(
self.model, gene_dataset_seq, train_size, test_size, seed)
self.train_fish, self.test_fish = self.train_test(
self.model, gene_dataset_fish, train_size, test_size, seed,
FishPosterior)
self.test_seq.to_monitor = ["reconstruction_error"]
self.test_fish.to_monitor = ["reconstruction_error"]
def train(self, n_epochs=20, lr=1e-3, weight_decay=1e-6, params=None):
self.adversarial_cls = Classifier(self.model.n_latent,
n_labels=self.model.n_batch,
n_layers=3)
if self.use_cuda:
self.adversarial_cls.cuda()
self.optimizer_cls = torch.optim.Adam(
filter(lambda p: p.requires_grad,
self.adversarial_cls.parameters()),
lr=lr,
weight_decay=weight_decay,
)
super().train(n_epochs=n_epochs, lr=1e-3, params=None)
@property
def posteriors_loop(self):
return ["train_seq", "train_fish"]
def loss(self, tensors_seq, tensors_fish):
sample_batch, local_l_mean, local_l_var, batch_index, labels = tensors_seq
reconst_loss, kl_divergence = self.model(
sample_batch,
local_l_mean,
local_l_var,
batch_index,
mode="scRNA",
weighting=self.weighting,
)
# If we want to add a classification loss
if self.cl_ratio != 0:
reconst_loss += self.cl_ratio * F.cross_entropy(
self.model.classify(sample_batch, mode="scRNA"),
labels.view(-1))
loss = torch.mean(reconst_loss + self.kl_weight * kl_divergence)
if (len(tensors_fish) == 7
): # depending on whether or not we have spatial coordinates
sample_batch_fish, local_l_mean, local_l_var, batch_index_fish, _, _, _ = (
tensors_fish)
else:
sample_batch_fish, local_l_mean, local_l_var, batch_index_fish, _ = (
tensors_fish)
reconst_loss_fish, kl_divergence_fish = self.model(
sample_batch_fish,
local_l_mean,
local_l_var,
batch_index_fish,
mode="smFISH",
)
loss_fish = torch.mean(reconst_loss_fish +
self.kl_weight * kl_divergence_fish)
loss = loss * sample_batch.size(
0) + loss_fish * sample_batch_fish.size(0)
loss /= sample_batch.size(0) + sample_batch_fish.size(0)
if self.epoch > self.warm_up:
sample_batch, local_l_mean, local_l_var, batch_index, labels = tensors_seq
z = self.model.sample_from_posterior_z(sample_batch, mode="scRNA")
cls_loss = self.scale * F.cross_entropy(
self.adversarial_cls(z),
torch.zeros_like(batch_index).view(-1))
if (len(tensors_fish) == 7
): # depending on whether or not we have spatial coordinates
sample_batch_fish, local_l_mean, local_l_var, batch_index_fish, _, _, _ = (
tensors_fish)
else:
sample_batch_fish, local_l_mean, local_l_var, batch_index_fish, _ = (
tensors_fish)
z = self.model.sample_from_posterior_z(sample_batch, mode="smFISH")
cls_loss += self.scale * F.cross_entropy(
self.adversarial_cls(z),
torch.ones_like(batch_index).view(-1))
self.optimizer_cls.zero_grad()
cls_loss.backward(retain_graph=True)
self.optimizer_cls.step()
else:
cls_loss = 0
return loss + loss_fish - cls_loss
def on_epoch_begin(self):
self.weighting = min(1, self.epoch / self.n_epochs_even)
self.kl_weight = (self.kl if self.kl is not None else min(
1, self.epoch / self.n_epochs_kl))
self.classification_ponderation = min(1, self.epoch / self.n_epochs_cl)
def test_cortex(save_path):
cortex_dataset = CortexDataset(save_path=save_path)
vae = VAE(cortex_dataset.nb_genes, cortex_dataset.n_batches)
trainer_cortex_vae = UnsupervisedTrainer(
vae, cortex_dataset, train_size=0.5, use_cuda=use_cuda
)
trainer_cortex_vae.train(n_epochs=1)
trainer_cortex_vae.train_set.reconstruction_error()
trainer_cortex_vae.train_set.differential_expression_stats()
trainer_cortex_vae.train_set.generate_feature_correlation_matrix(
n_samples=2, correlation_type="pearson"
)
trainer_cortex_vae.train_set.generate_feature_correlation_matrix(
n_samples=2, correlation_type="spearman"
)
trainer_cortex_vae.train_set.imputation(n_samples=1)
trainer_cortex_vae.test_set.imputation(n_samples=5)
trainer_cortex_vae.corrupt_posteriors(corruption="binomial")
trainer_cortex_vae.corrupt_posteriors()
trainer_cortex_vae.train(n_epochs=1)
trainer_cortex_vae.uncorrupt_posteriors()
trainer_cortex_vae.train_set.imputation_benchmark(
n_samples=1, show_plot=False, title_plot="imputation", save_path=save_path
)
trainer_cortex_vae.train_set.generate_parameters()
n_cells, n_genes = (
len(trainer_cortex_vae.train_set.indices),
cortex_dataset.nb_genes,
)
n_samples = 3
(dropout, means, dispersions,) = trainer_cortex_vae.train_set.generate_parameters()
assert dropout.shape == (n_cells, n_genes) and means.shape == (n_cells, n_genes)
assert dispersions.shape == (n_cells, n_genes)
(dropout, means, dispersions,) = trainer_cortex_vae.train_set.generate_parameters(
n_samples=n_samples
)
assert dropout.shape == (n_samples, n_cells, n_genes)
assert means.shape == (n_samples, n_cells, n_genes,)
(dropout, means, dispersions,) = trainer_cortex_vae.train_set.generate_parameters(
n_samples=n_samples, give_mean=True
)
assert dropout.shape == (n_cells, n_genes) and means.shape == (n_cells, n_genes)
full = trainer_cortex_vae.create_posterior(
vae, cortex_dataset, indices=np.arange(len(cortex_dataset))
)
x_new, x_old = full.generate(n_samples=10)
assert x_new.shape == (cortex_dataset.nb_cells, cortex_dataset.nb_genes, 10)
assert x_old.shape == (cortex_dataset.nb_cells, cortex_dataset.nb_genes)
trainer_cortex_vae.train_set.imputation_benchmark(
n_samples=1, show_plot=False, title_plot="imputation", save_path=save_path
)
svaec = SCANVI(
cortex_dataset.nb_genes, cortex_dataset.n_batches, cortex_dataset.n_labels
)
trainer_cortex_svaec = JointSemiSupervisedTrainer(
svaec, cortex_dataset, n_labelled_samples_per_class=3, use_cuda=use_cuda
)
trainer_cortex_svaec.train(n_epochs=1)
trainer_cortex_svaec.labelled_set.accuracy()
trainer_cortex_svaec.full_dataset.reconstruction_error()
svaec = SCANVI(
cortex_dataset.nb_genes, cortex_dataset.n_batches, cortex_dataset.n_labels
)
trainer_cortex_svaec = AlternateSemiSupervisedTrainer(
svaec, cortex_dataset, n_labelled_samples_per_class=3, use_cuda=use_cuda
)
trainer_cortex_svaec.train(n_epochs=1, lr=1e-2)
trainer_cortex_svaec.unlabelled_set.accuracy()
data_train, labels_train = trainer_cortex_svaec.labelled_set.raw_data()
data_test, labels_test = trainer_cortex_svaec.unlabelled_set.raw_data()
compute_accuracy_svc(
data_train,
labels_train,
data_test,
labels_test,
param_grid=[{"C": [1], "kernel": ["linear"]}],
)
compute_accuracy_rf(
data_train,
labels_train,
data_test,
labels_test,
param_grid=[{"max_depth": [3], "n_estimators": [10]}],
)
cls = Classifier(cortex_dataset.nb_genes, n_labels=cortex_dataset.n_labels)
cls_trainer = ClassifierTrainer(cls, cortex_dataset)
cls_trainer.train(n_epochs=1)
cls_trainer.train_set.accuracy()
class SCANVI(VAE):
r"""A semi-supervised Variational auto-encoder model - inspired from M1 + M2 model,
as described in (https://arxiv.org/pdf/1406.5298.pdf). SCANVI stands for single-cell annotation using
variational inference.
:param n_input: Number of input genes
:param n_batch: Number of batches
:param n_labels: Number of labels
:param n_hidden: Number of nodes per hidden layer
:param n_latent: Dimensionality of the latent space
:param n_layers: Number of hidden layers used for encoder and decoder NNs
:param dropout_rate: Dropout rate for neural networks
:param dispersion: One of the following
* ``'gene'`` - dispersion parameter of NB is constant per gene across cells
* ``'gene-batch'`` - dispersion can differ between different batches
* ``'gene-label'`` - dispersion can differ between different labels
* ``'gene-cell'`` - dispersion can differ for every gene in every cell
:param log_variational: Log variational distribution
:param reconstruction_loss: One of
* ``'nb'`` - Negative binomial distribution
* ``'zinb'`` - Zero-inflated negative binomial distribution
:param y_prior: If None, initialized to uniform probability over cell types
:param labels_groups: Label group designations
:param use_labels_groups: Whether to use the label groups
Examples:
>>> gene_dataset = CortexDataset()
>>> scanvi = SCANVI(gene_dataset.nb_genes, n_batch=gene_dataset.n_batches * False,
... n_labels=gene_dataset.n_labels)
>>> gene_dataset = SyntheticDataset(n_labels=3)
>>> scanvi = SCANVI(gene_dataset.nb_genes, n_batch=gene_dataset.n_batches * False,
... n_labels=3, y_prior=torch.tensor([[0.1,0.5,0.4]]), labels_groups=[0,0,1])
"""
def __init__(self,
n_input: int,
n_batch: int = 0,
n_labels: int = 0,
n_hidden: int = 128,
n_latent: int = 10,
n_layers: int = 1,
dropout_rate: float = 0.1,
dispersion: str = "gene",
log_variational: bool = True,
reconstruction_loss: str = "zinb",
y_prior=None,
labels_groups: Sequence[int] = None,
use_labels_groups: bool = False,
classifier_parameters: dict = dict()):
super(SCANVI, self).__init__(n_input,
n_hidden=n_hidden,
n_latent=n_latent,
n_layers=n_layers,
dropout_rate=dropout_rate,
n_batch=n_batch,
dispersion=dispersion,
log_variational=log_variational,
reconstruction_loss=reconstruction_loss)
self.n_labels = n_labels
self.n_latent_layers = 2
# Classifier takes n_latent as input
cls_parameters = {
"n_layers": n_layers,
"n_hidden": n_hidden,
"dropout_rate": dropout_rate
}
cls_parameters.update(classifier_parameters)
self.classifier = Classifier(n_latent,
n_labels=self.n_labels,
**cls_parameters)
self.encoder_z2_z1 = Encoder(n_latent,
n_latent,
n_cat_list=[self.n_labels],
n_layers=n_layers,
n_hidden=n_hidden,
dropout_rate=dropout_rate)
self.decoder_z1_z2 = Decoder(n_latent,
n_latent,
n_cat_list=[self.n_labels],
n_layers=n_layers,
n_hidden=n_hidden,
dropout_rate=dropout_rate)
self.y_prior = torch.nn.Parameter(y_prior if y_prior is not None else
(1 / n_labels) *
torch.ones(1, n_labels),
requires_grad=False)
self.use_labels_groups = use_labels_groups
self.labels_groups = np.array(
labels_groups) if labels_groups is not None else None
if self.use_labels_groups:
assert labels_groups is not None, "Specify label groups"
unique_groups = np.unique(self.labels_groups)
self.n_groups = len(unique_groups)
assert (unique_groups == np.arange(self.n_groups)).all()
self.classifier_groups = Classifier(n_latent, n_hidden,
self.n_groups, n_layers,
dropout_rate)
self.groups_index = torch.nn.ParameterList([
torch.nn.Parameter(torch.tensor(
(self.labels_groups == i).astype(np.uint8),
dtype=torch.uint8),
requires_grad=False)
for i in range(self.n_groups)
])
def train(self, mode=True):
super(SCANVI, self).train(mode=mode)
self.classifier.train(mode=mode, batch_norm=True)
def classify(self, x):
z, _, _ = self.z_encoder(
torch.log(1 + x)) # not using the sampled version, here z = qz_m
if self.use_labels_groups:
w_g = self.classifier_groups(z)
unw_y = self.classifier(z)
w_y = torch.zeros_like(unw_y)
for i, group_index in enumerate(self.groups_index):
unw_y_g = unw_y[:, group_index]
w_y[:, group_index] = unw_y_g / (
unw_y_g.sum(dim=-1, keepdim=True) + 1e-8)
w_y[:, group_index] *= w_g[:, [i]]
else:
w_y = self.classifier(z)
return w_y
def get_latents(self, x, y=None):
zs = super(SCANVI, self).get_latents(x)
qz2_m, qz2_v, z2 = self.encoder_z2_z1(zs[0], y)
if not self.training:
z2 = qz2_m
return [zs[0], z2]
def forward(self, x, local_l_mean, local_l_var, batch_index=None, y=None):
is_labelled = False if y is None else True
x_ = torch.log(1 + x)
qz1_m, qz1_v, z1 = self.z_encoder(x_)
ql_m, ql_v, library = self.l_encoder(x_)
# Enumerate choices of label
ys, z1s = (broadcast_labels(y, z1, n_broadcast=self.n_labels))
qz2_m, qz2_v, z2 = self.encoder_z2_z1(z1s, ys)
pz1_m, pz1_v = self.decoder_z1_z2(z2, ys)
px_scale, px_r, px_rate, px_dropout = self.decoder(
self.dispersion, z1, library, batch_index)
reconst_loss = self._reconstruction_loss(x, px_rate, px_r, px_dropout,
batch_index, y)
# KL Divergence
mean = torch.zeros_like(qz2_m)
scale = torch.ones_like(qz2_v)
kl_divergence_z2 = kl(Normal(qz2_m, torch.sqrt(qz2_v)),
Normal(mean, scale)).sum(dim=1)
loss_z1_unweight = -Normal(pz1_m,
torch.sqrt(pz1_v)).log_prob(z1s).sum(dim=-1)
loss_z1_weight = Normal(qz1_m,
torch.sqrt(qz1_v)).log_prob(z1).sum(dim=-1)
kl_divergence_l = kl(Normal(ql_m, torch.sqrt(ql_v)),
Normal(local_l_mean,
torch.sqrt(local_l_var))).sum(dim=1)
if is_labelled:
return reconst_loss + loss_z1_weight + loss_z1_unweight, kl_divergence_z2 + kl_divergence_l
probs = self.classifier(z1)
reconst_loss += (loss_z1_weight + (
(loss_z1_unweight).view(self.n_labels, -1).t() * probs).sum(dim=1))
kl_divergence = (kl_divergence_z2.view(self.n_labels, -1).t() *
probs).sum(dim=1)
kl_divergence += kl(
Categorical(probs=probs),
Categorical(probs=self.y_prior.repeat(probs.size(0), 1)))
kl_divergence += kl_divergence_l
return reconst_loss, kl_divergence