Finetune tutorial (#30)

* added finetuning tutorial

* finetuning tutorial

* fixed wandb connection

* ran test

* added dataset

Author: avantikalal

src/decima/cli/finetune.py

@@ -16,6 +16,12 @@
     type=str,
     help="Model path or replication number. If a path is provided, the model will be loaded from the path. If a replication number is provided, the model will be loaded from the replication number.",
 )
+@click.option(
+    "--device",
+    type=str,
+    default="0",
+    help="Device to use. Default: 0",
+)
 @click.option("--matrix-file", required=True, help="Matrix file path.")
 @click.option("--h5-file", required=True, help="H5 file path.")
 @click.option("--outdir", required=True, help="Output directory path to save model checkpoints.")
@@ -33,6 +39,7 @@
 def cli_finetune(
     name,
     model,
+    device,
     matrix_file,
     h5_file,
     outdir,
@@ -66,11 +73,14 @@ def cli_finetune(
     )
     val_dataset = HDF5Dataset(h5_file=h5_file, ad=ad, key="val", max_seq_shift=0)
 
+    if isinstance(device, str) and device.isdigit():
+        device = int(device)
+        
     train_params = {
         "name": name,
         "batch_size": batch_size,
         "num_workers": num_workers,
-        "devices": 0,
+        "devices": device,
         "logger": train_logger,
         "save_dir": outdir,
         "max_epochs": epochs,
@@ -94,10 +104,12 @@ def cli_finetune(
     logger.info("Initializing model")
     model = LightningModel(model_params=model_params, train_params=train_params)
 
-    logger.info("Training")
-    if logger == "wandb":
-        wandb.login(host="https://genentech.wandb.io")
+    if train_logger == "wandb":
+        logger.info("Connecting to wandb.")
+        wandb.login(host="https://genentech.wandb.io", anonymous="never")
         run = wandb.init(project="decima", dir=name, name=name)
+
+    logger.info("Training")
     model.train_on_dataset(train_dataset, val_dataset)
     train_dataset.close()
     val_dataset.close()

src/decima/data/write_hdf5.py

@@ -4,7 +4,7 @@
 from grelu.sequence.utils import get_unique_length
 
 
-def write_hdf5(file, ad, pad=0):
+def write_hdf5(file, ad, pad=0, genome="hg38"):
     # Calculate seq_len
     seq_len = get_unique_length(ad.var)
 
@@ -45,7 +45,7 @@ def write_hdf5(file, ad, pad=0):
         arr = ad.var[["chrom", "start", "end", "strand"]].copy()
         arr.start = arr.start - pad
         arr.end = arr.end + pad
-        arr = convert_input_type(arr, "indices", genome="hg38")
+        arr = convert_input_type(arr, "indices", genome=genome)
         print(f"Writing sequence array of shape: {arr.shape}")
         f.create_dataset("sequences", shape=arr.shape, dtype=np.int8, data=arr)
 

tutorials/2-variant-effect-prediction.ipynb

@@ -4859,7 +4859,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "decima",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -4873,9 +4873,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.12"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }