def trainSCANVI(gene_dataset,
model_type,
filename,
rep,
nlayers=2,
reconstruction_loss: str = "zinb"):
vae_posterior = trainVAE(gene_dataset,
filename,
rep,
reconstruction_loss=reconstruction_loss)
filename = '../' + filename + '/' + model_type + '.' + reconstruction_loss + '.rep' + str(
rep) + '.pkl'
scanvi = SCANVI(gene_dataset.nb_genes,
gene_dataset.n_batches,
gene_dataset.n_labels,
n_layers=nlayers,
reconstruction_loss=reconstruction_loss)
scanvi.load_state_dict(vae_posterior.model.state_dict(), strict=False)
if model_type == 'scanvi1':
trainer_scanvi = AlternateSemiSupervisedTrainer(
scanvi,
gene_dataset,
classification_ratio=0,
n_epochs_classifier=100,
lr_classification=5 * 1e-3)
labelled = np.where(gene_dataset.batch_indices.ravel() == 0)[0]
labelled = np.random.choice(labelled, len(labelled), replace=False)
trainer_scanvi.labelled_set = trainer_scanvi.create_posterior(
indices=labelled)
trainer_scanvi.unlabelled_set = trainer_scanvi.create_posterior(
indices=(gene_dataset.batch_indices.ravel() == 1))
elif model_type == 'scanvi2':
trainer_scanvi = AlternateSemiSupervisedTrainer(
scanvi,
gene_dataset,
classification_ratio=0,
n_epochs_classifier=100,
lr_classification=5 * 1e-3)
labelled = np.where(gene_dataset.batch_indices.ravel() == 1)[0]
labelled = np.random.choice(labelled, len(labelled), replace=False)
trainer_scanvi.labelled_set = trainer_scanvi.create_posterior(
indices=labelled)
trainer_scanvi.unlabelled_set = trainer_scanvi.create_posterior(
indices=(gene_dataset.batch_indices.ravel() == 0))
elif model_type == 'scanvi0':
trainer_scanvi = SemiSupervisedTrainer(scanvi,
gene_dataset,
classification_ratio=0,
n_epochs_classifier=100,
lr_classification=5 * 1e-3)
trainer_scanvi.labelled_set = trainer_scanvi.create_posterior(
indices=(gene_dataset.batch_indices.ravel() < 0))
trainer_scanvi.unlabelled_set = trainer_scanvi.create_posterior(
indices=(gene_dataset.batch_indices.ravel() >= 0))
else:
trainer_scanvi = SemiSupervisedTrainer(scanvi,
gene_dataset,
classification_ratio=10,
n_epochs_classifier=100,
lr_classification=5 * 1e-3)
if os.path.isfile(filename):
trainer_scanvi.model.load_state_dict(torch.load(filename))
trainer_scanvi.model.eval()
else:
trainer_scanvi.train(n_epochs=5)
torch.save(trainer_scanvi.model.state_dict(), filename)
return trainer_scanvi
