update step3 imputation

Author: xingzhongyu

test_automl/step2_imputation_deepimpute.py

@@ -58,7 +58,7 @@ def train(config):
     print('deepimpute')
     print(f'rmses: {rmses}')
     print(f'rmses: {np.mean(rmses)} +/- {np.std(rmses)}')
-    return ({"rmses": np.mean(rmses)})
+    return ({"scores": np.mean(rmses)})
 
 
 def startSweep(parameters_dict) -> Tuple[Dict[str, Any], Callable[..., Any]]:
@@ -111,7 +111,7 @@ def startSweep(parameters_dict) -> Tuple[Dict[str, Any], Callable[..., Any]]:
     })
     sweep_config = {'method': 'grid'}
     sweep_config['parameters'] = parameters_dict
-    metric = {'name': 'rmses', 'goal': 'minimize'}
+    metric = {'name': 'scores', 'goal': 'minimize'}
 
     sweep_config['metric'] = metric
     return sweep_config, train  #Return function configuration and training function

test_automl/step3_config.py

@@ -3,13 +3,14 @@
 
 import optuna
 import scanpy as sc
-import wandb
 from fun2code import fun2code_dict
 from optuna.integration.wandb import WeightsAndBiasesCallback
 from step2_config import pipline2fun_dict
 
+import wandb
 from dance.transforms.cell_feature import CellPCA, CellSVD, WeightedFeaturePCA
 from dance.transforms.filter import FilterGenesPercentile, FilterGenesRegression
+from dance.transforms.gene_holdout import GeneHoldout
 from dance.transforms.interface import AnnDataTransform
 from dance.transforms.mask import CellwiseMaskData, MaskData
 from dance.transforms.misc import Compose, SetConfig
@@ -144,6 +145,12 @@ def mask_data(method_name: str, trial: optuna.Trial):
     return MaskData(mask_rate=trial.suggest_float(method_name + "mask_rate", 0.01, 0.5))
 
 
+@set_method_name
+def gene_hold_out(method_name: str, trial: optuna.Trial):
+    return GeneHoldout(n_top=trial.suggest_int(method_name + "n_top", 1, 10),
+                       batch_size=trial.suggest_categorical(method_name + "batch_size", [256, 512, 1024]))
+
+
 # # 获取当前文件中的所有函数
 # functions = [(name,obj) for name, obj in inspect.getmembers(
 #     sys.modules[__name__]) if inspect.isfunction(obj)]
@@ -197,13 +204,13 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
-def get_optimizer(project, objective, n_trials=2):
+def get_optimizer(project, objective, n_trials=2, direction="maximize"):
     """Get optimizer."""
     wandb_kwargs = {"project": project}
     wandbc = WeightsAndBiasesCallback(wandb_kwargs=wandb_kwargs, as_multirun=True)
     decorator = log_in_wandb(wandbc)
     decorator_function = decorator(objective)
-    study = optuna.create_study()
+    study = optuna.create_study(direction=direction)
 
     def wrapper():
         study.optimize(decorator_function, n_trials=n_trials, callbacks=[wandbc])

test_automl/step3_imputation_deepimpute.py

@@ -0,0 +1,87 @@
+import torch
+from step3_config import get_optimizer, get_transforms
+
+from dance import logger
+from dance.datasets.singlemodality import ImputationDataset
+from dance.modules.single_modality.imputation.deepimpute import DeepImpute
+from dance.registry import DotDict
+from dance.transforms.misc import Compose, SetConfig
+from dance.utils import set_seed
+
+fun_list = ["filter_gene_by_count", "filter_cell_by_count", "log1p", "gene_hold_out", "cell_wise_mask_data"]
+device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
+import numpy as np
+
+
+def objective(trial):
+    parameters_dict = {
+        'dropout': 0.1,
+        'lr': 1e-5,
+        'n_epochs': 5,
+        'batch_size': 64,
+        'sub_outputdim': 512,
+        'hidden_dim': 256,
+        'patience': 20,
+        'min_cells': 0.05,
+        "n_top": 5,
+        "train_size": 0.9,
+        "mask_rate": 0.1,
+        "cache": False,
+        "mask": True,  #避免出现与超参数流程重复的情况，一般没有
+        "seed": 0,
+        "num_runs": 1
+    }
+    parameters_config = {}
+    parameters_config.update(parameters_dict)
+    parameters_config = DotDict(parameters_config)
+    rmses = []
+    for seed in range(parameters_config.seed, parameters_config.seed + parameters_config.num_runs):
+        set_seed(seed)
+        dataset = "mouse_brain_data"
+        data_dir = "./test_automl/data"
+        dataloader = ImputationDataset(data_dir=data_dir, dataset=dataset, train_size=parameters_config.train_size)
+        # preprocessing_pipeline = DeepImpute.preprocessing_pipeline(min_cells=parameters_config.min_cells, n_top=parameters_config.n_top,
+        #                                                            sub_outputdim=parameters_config.sub_outputdim, mask=parameters_config.mask,
+        #                                                            seed=seed, mask_rate=parameters_config.mask_rate)
+        transforms = get_transforms(trial=trial, fun_list=fun_list, set_data_config=False, save_raw=True)
+        if transforms is None:
+            logger.warning("skip transforms")
+            return {"scores": 0}
+        transforms.append(
+            SetConfig({
+                "feature_channel": [None, None, "targets", "predictors", "train_mask"],
+                "feature_channel_type": ["X", "raw_X", "uns", "uns", "layers"],
+                "label_channel": [None, None],
+                "label_channel_type": ["X", "raw_X"],
+            }))
+        preprocessing_pipeline = Compose(*transforms, log_level="INFO")
+        data = dataloader.load_data(transform=preprocessing_pipeline, cache=parameters_config.cache)
+
+        if parameters_config.mask:
+            X, X_raw, targets, predictors, mask = data.get_x(return_type="default")
+        else:
+            mask = None
+            X, X_raw, targets, predictors = data.get_x(return_type="default")
+        X = torch.tensor(X.toarray()).float()
+        X_raw = torch.tensor(X_raw.toarray()).float()
+        X_train = X * mask
+        model = DeepImpute(predictors, targets, dataset, parameters_config.sub_outputdim, parameters_config.hidden_dim,
+                           parameters_config.dropout, seed, 1)
+
+        model.fit(X_train, X_train, mask, parameters_config.batch_size, parameters_config.lr,
+                  parameters_config.n_epochs, parameters_config.patience)
+        imputed_data = model.predict(X_train, mask)
+        score = model.score(X, imputed_data, mask, metric='RMSE')
+        print("RMSE: %.4f" % score)
+        rmses.append(score)
+
+    print('deepimpute')
+    print(f'rmses: {rmses}')
+    print(f'rmses: {np.mean(rmses)} +/- {np.std(rmses)}')
+    return ({"scores": np.mean(rmses)})
+
+
+if __name__ == "__main__":
+    start_optimizer = get_optimizer(project="step3-imputation-deepimpute-project", objective=objective, n_trials=10,
+                                    direction="minimize")
+    start_optimizer()