diff --git a/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml index 27151ff6..b3469f38 100644 --- a/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml +++ b/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml @@ -18,12 +18,13 @@ engines: image: openproblems/base_pytorch_nvidia:1 setup: - type: docker - run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git + run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@be8bee65c146316e579ae092c43247284ba801fe - type: python packages: - lightgbm>=4.0 - tensorflow>=2.12 - scikit-learn>=1.1 + - joblib>=1.4 - mudata>=0.2 - muon>=0.1 - fast-array-utils diff --git a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml index 46642bb3..ef7ec762 100644 --- a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml +++ b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml @@ -17,16 +17,27 @@ arguments: - name: "--n_folds" type: integer default: 5 - description: Number of cross-validation folds for LightGBM and neural network training. + description: Number of cross-validation folds for the neural networks (one network per fold, predictions averaged). + - name: "--lgbm_n_folds" + type: integer + default: 3 + description: | + Number of cross-validation folds for the LightGBM models, whose out-of-fold predictions are + features of the neural networks. Fewer folds cut the LightGBM time proportionally at a small + cost in the quality of these features. - name: "--lgbm_boost_rounds" type: integer - default: 10000 - description: Maximum LightGBM boosting rounds (early stopping applies). + default: 100 + description: | + Maximum LightGBM boosting rounds (early stopping applies). The original solution used up to + 10000 rounds at learning rate 0.01, which takes days on the benchmark datasets. At learning + rate 0.1, 100 rounds keep the LightGBM per-target quality within 0.015 (NeurIPS 2021) / 0.003 + (NeurIPS 2022) of 300 rounds, and the final predictions of the networks are unchanged. info: test_default: 50 - name: "--lgbm_early_stopping" type: integer - default: 100 + default: 20 description: LightGBM early stopping patience (rounds without improvement). info: test_default: 10 @@ -40,6 +51,26 @@ arguments: type: integer default: 100 description: TSVD components for reducing LightGBM predictions before NN input. + - name: "--lgbm_learning_rate" + type: double + default: 0.1 + description: | + LightGBM learning rate. The original solution used 0.01 with up to 10000 rounds; 0.1 with a few + hundred rounds reaches the same LightGBM quality in a fraction of the time. + - name: "--lgbm_max_bin" + type: integer + default: 63 + description: | + Number of histogram bins per feature in LightGBM (library default 255). 63 halves the memory + of the binned whole-transcriptome inputs held by every worker process (3 instead of 6.6 GB per + fold on the NeurIPS 2022 CITE data) at identical validation loss and slightly faster rounds. + - name: "--lgbm_n_jobs" + type: integer + default: -1 + description: | + Number of worker processes training different protein targets in parallel; the allocated CPUs + are divided among them. -1 uses one process per allocated CPU (each with one thread), which is + several times faster than one target at a time with all threads. engines: - type: docker image: nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04 @@ -60,13 +91,14 @@ engines: - jsonschema - lightgbm>=4.0 - scikit-learn>=1.1 + - joblib>=1.4 - mudata>=0.2 - muon>=0.1 - fast-array-utils github: - openproblems-bio/core#subdirectory=packages/python/openproblems - type: docker - run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git + run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@be8bee65c146316e579ae092c43247284ba801fe runners: - type: executable - type: nextflow diff --git a/src/methods/senkin_tmp/senkin_tmp_train/script.py b/src/methods/senkin_tmp/senkin_tmp_train/script.py index 0c6c1eaf..bc2c87e4 100644 --- a/src/methods/senkin_tmp/senkin_tmp_train/script.py +++ b/src/methods/senkin_tmp/senkin_tmp_train/script.py @@ -29,12 +29,16 @@ "input_test_mod1": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal/test_mod1.h5ad", "output": "output_model.pkl", "n_folds": 5, - "lgbm_boost_rounds": 10000, - "lgbm_early_stopping": 100, + "lgbm_n_folds": 3, + "lgbm_boost_rounds": 100, + "lgbm_early_stopping": 20, "nn_epochs": 100, "n_tsvd_components": 100, + "lgbm_learning_rate": 0.1, + "lgbm_max_bin": 63, + "lgbm_n_jobs": -1, } -meta = {"name": "senkin_tmp", "resources_dir": "src/methods/senkin_tmp/senkin_tmp_train", "cpus": None} +meta = {"name": "senkin_tmp", "resources_dir": "src/methods/senkin_tmp/senkin_tmp_train", "cpus": None, "memory_gb": None} ## VIASH END sys.path.append(meta["resources_dir"]) @@ -148,6 +152,7 @@ def _feature_names(adata): Y_prot_raw = to_dense(adata_prot_train.layers.get("counts", adata_prot_train.X), dtype=np.float64) folds = KFold(n_splits=par["n_folds"], shuffle=True, random_state=666) +lgbm_folds = KFold(n_splits=par["lgbm_n_folds"], shuffle=True, random_state=666) n_tsvd = par["n_tsvd_components"] boost_rounds = par["lgbm_boost_rounds"] early_stop = par["lgbm_early_stopping"] @@ -157,9 +162,14 @@ def _feature_names(adata): # meta["cpus"], so the threads oversubscribe and thrash -- the same class of slowdown # fixed for guanlab in #59. Leave the library default when cpus is unknown (local runs). _n_threads = meta.get("cpus") -if _n_threads: - for _p in (lgbm_params_1, lgbm_params_2, lgbm_params_3, lgbm_params_4): +for _p in (lgbm_params_1, lgbm_params_2, lgbm_params_3, lgbm_params_4): + if _n_threads: _p["num_threads"] = _n_threads + _p["learning_rate"] = par["lgbm_learning_rate"] + _p["max_bin"] = par["lgbm_max_bin"] +lgbm_n_jobs = par["lgbm_n_jobs"] +# Let the library keep the LightGBM worker processes within the allocated memory (a tenth is left for the rest). +lgbm_memory_budget_gb = meta["memory_gb"] * 0.9 if meta.get("memory_gb") else None # --------------------------------------------------------------------------- # LightGBM — 4 models, train+test passed together (original design) @@ -168,11 +178,14 @@ def _feature_names(adata): def _lgbm(X_all, Y, params, description): logger.info(f"Training LightGBM {description}...") return get_lgbm_predictions( - X_all[train_idx], Y, X_all[test_idx], folds, params, + X_all[train_idx], Y, X_all[test_idx], lgbm_folds, params, n_tsvd_components=n_tsvd, num_boost_round=boost_rounds, early_stopping_rounds=early_stop, + n_jobs=lgbm_n_jobs, memory_budget_gb=lgbm_memory_budget_gb, ) lgbm1_svd_all = _lgbm(X_lognorm_all, Y_prot_train, lgbm_params_1, "model 1 (log-normalized RNA -> proteins)") +del X_lognorm_all # not needed any more; keeps the parent's footprint (and hence the memory left for workers) small +gc.collect() X_comb_all = np.concatenate([X_clr_tsvd_all, X_raw_selected_all, X_sqrt_tsvd_all, X_sqrt_pca_all], axis=1) lgbm2_svd_all = _lgbm(X_comb_all, Y_prot_train, lgbm_params_2, "model 2 (CLR-TSVD + selected genes + normalized TSVD/PCA -> proteins)") @@ -180,7 +193,7 @@ def _lgbm(X_all, Y, params, description): lgbm3_svd_all = _lgbm(X_counts_all, Y_prot_train, lgbm_params_3, "model 3 (raw counts -> proteins)") lgbm4_svd_all = _lgbm(X_counts_all, Y_prot_raw, lgbm_params_4, "model 4 (raw counts -> raw proteins)") -del X_counts_all, X_lognorm_all +del X_counts_all gc.collect() # ---------------------------------------------------------------------------